From 5d07d789179050f95d961660498709e9b432ee17 Mon Sep 17 00:00:00 2001
From: Melissa Linkert <melissa@glencoesoftware.com>
Date: Thu, 5 Sep 2024 20:36:56 -0500
Subject: [PATCH] Add `--chunk` option to allow re-chunking when writing v3

---
 .../com/glencoesoftware/zarr/Convert.java     | 19 ++++++
 .../zarr/test/ConversionTest.java             | 59 +++++++++++++++++++
 2 files changed, 78 insertions(+)

diff --git a/src/main/java/com/glencoesoftware/zarr/Convert.java b/src/main/java/com/glencoesoftware/zarr/Convert.java
index f10b53d..aeafcc7 100644
--- a/src/main/java/com/glencoesoftware/zarr/Convert.java
+++ b/src/main/java/com/glencoesoftware/zarr/Convert.java
@@ -68,6 +68,7 @@ public class Convert implements Callable<Integer> {
 
   private ShardConfiguration shardConfig;
   private int[] requestedShard; // the requested size for custom sharding
+  private int[] requestedChunkSize;
   private String[] codecs;
 
   /**
@@ -123,6 +124,21 @@ public void setWriteV2(boolean v2) {
     writeV2 = v2;
   }
 
+  @Option(
+    names = "--chunk",
+    description = "'t,c,z,y,x' (optional comma-separated custom shard size)",
+    defaultValue = ""
+  )
+  public void setChunk(String chunk) {
+    if (chunk != null && !chunk.isEmpty()) {
+      String[] chunkSize = chunk.split(",");
+      requestedChunkSize = new int[chunkSize.length];
+      for (int i=0; i<chunkSize.length; i++) {
+        requestedChunkSize[i] = Integer.parseInt(chunkSize[i]);
+      }
+    }
+  }
+
   @Option(
     names = "--shard",
     description = "'single' (one shard per array), " +
@@ -284,6 +300,9 @@ public void convertToV3() throws Exception {
             ZarrArray tile = field.openArray("/" + res);
             LOGGER.info("opened array {}", resolutionPath);
             int[] originalChunkSizes = tile.getChunks();
+            if (requestedChunkSize != null) {
+              originalChunkSizes = requestedChunkSize;
+            }
             int[] shape = tile.getShape();
 
             int[] chunkSizes = new int[originalChunkSizes.length];
diff --git a/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java b/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java
index 080df85..50c6031 100644
--- a/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java
+++ b/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java
@@ -316,6 +316,65 @@ public void testCodecs() throws Exception {
     }
   }
 
+  /**
+   * Test different chunk sizes without sharding.
+   */
+  @Test
+  public void testChunkSizes() throws Exception {
+    input = fake("sizeX", "5120", "sizeY", "5120", "sizeZ", "4");
+    assertBioFormats2Raw();
+
+    String[] chunkOptions = new String[] {
+      "1,1,2,32,32",
+      "1,1,4,512,512",
+      "1,1,1,2048,2048"
+    };
+    int[][] chunkSizes = new int[][] {
+      {1, 1, 2, 32, 32},
+      {1, 1, 4, 512, 512},
+      {1, 1, 1, 2048, 2048}
+    };
+
+    for (int opt=0; opt<chunkOptions.length; opt++) {
+      // first convert v2 produced by bioformats2raw to v3
+      Path v3Output = tmp.newFolder().toPath().resolve("v3-test");
+      Convert v3Converter = new Convert();
+      v3Converter.setInput(output.toString());
+      v3Converter.setOutput(v3Output.toString());
+
+      v3Converter.setChunk(chunkOptions[opt]);
+      v3Converter.convertToV3();
+
+      Store store = new FilesystemStore(v3Output);
+      Array resolution = Array.open(store.resolve("0", "0"));
+
+      int[] chunkSize = chunkSizes[opt];
+      Assert.assertArrayEquals(resolution.metadata.chunkShape(), chunkSize);
+
+      // now convert v3 back to v2
+      Path roundtripOutput = tmp.newFolder().toPath().resolve("v2-roundtrip-test");
+      Convert v2Converter = new Convert();
+      v2Converter.setInput(v3Output.toString());
+      v2Converter.setOutput(roundtripOutput.toString());
+      v2Converter.setWriteV2(true);
+      v2Converter.convertToV2();
+
+      Path originalOMEXML = output.resolve("OME").resolve("METADATA.ome.xml");
+      Path roundtripOMEXML = roundtripOutput.resolve("OME").resolve("METADATA.ome.xml");
+
+      // make sure the OME-XML is present and not changed
+      Assert.assertEquals(Files.readAllLines(originalOMEXML), Files.readAllLines(roundtripOMEXML));
+
+      // since the image is small, make sure all pixels are identical in both resolutions
+      for (int r=0; r<6; r++) {
+        ZarrArray original = ZarrGroup.open(output.resolve("0")).openArray(String.valueOf(r));
+        ZarrArray roundtrip = ZarrGroup.open(roundtripOutput.resolve("0")).openArray(String.valueOf(r));
+
+        compareZarrArrays(original, roundtrip);
+      }
+    }
+  }
+
   /**
    * Test different default sharding options
    */