Feat: Add stitching multiple ome zarr files together

[jimmyjam100]

Added the ability for a user to stitch a bunch of ome zarr files together into a larger ome zarr. This PR largerly modifies and generalizes code that was used for spool.py

Also allows for the reading of zarr files from dandi

makes writing to zarr more dask friendly

I swear this PR is not that large its just poetry.lock needed to be rewritten but its auto generated so does not need to be reviewed

NOTES FOR AFTER SITE VISIT:
During the process of getting mode code running in time for the site visit I learned several things and ran into several difficulties. This is a log of all the problems I ran into, how I solved them, and how I would solve them differently if I take up this code again and am not under time crunch. The strict deadline from the site visit caused me to come to not super rational solutions to all of my code as you will soon see.

Most notably is that it seems when a zarr-python object is loaded into memory it stays in memory until the object is removed via garbage collection. This lead me to the solution of having two different objects for shapes and data so that garbage collection can easily delete the object containing data. This is far from an elegant solution, but it worked and did not have the time for constant revisions. If I were to redo this I would probably find a way to just load the data of the objects used to calculate shape, then manually remove all pointers to that object once it was done writing

After testing multiple different methods of skew correction it seems like using a wrapper worked the quickest. Like above this is far from the most elegant solution but I tested several ideas including just using dask_image.ndinterp.map_coordinates and dask_image.ndinterp.affine_transform with no wrapper. This wrapper works by mimicking the functionality of a dask array but also applying deskewing to it on the fly. This includes modifying the shape of the object. I actually like this solution. Having a wrapper that can edit the shape of an object is super convenient for down the line calculations. The major downside is it does take a whole object to do what could be done with a couple lines of code so it is messy.

I also ran into issues getting everything to run within 12 hours and ended up adding parameters to split runs along the x axis and combine them afterwards. This also had the benefit of loading only small parts of chunks into memory at a time when writing to a zarr file which was good because a single chunk would sometimes exceed the 128 GB memory limit for mit_normal. However after coming up with a checkpoint system and using mit_preemptable I think this addition is not needed. mit_preemptable would make it so I could use 512 GB of memory and the checkpoint system would mean I don't have to worry about time constraints. One possible upside of this method is that it allows for stitching multiple different sections in parallel across nodes which might be faster than how dask would do it (Have not tested). In the future I will probably just remove this feature to remove feature clutter especially because I think it does not work well with sharding.

I was also running into issues generating pyramids for very large ome zarrs because of OOM errors. My solution of only generating a pyramid one chunk at a time makes sense but the implementation is sloppy. Especially because in order to get it to work quickly I had to implement it in zarr-python as opposed to the parent zarr class.

In order to run this on the MIT cluster I would run the following scripts.

First I would preprocess stripe correction using

#!/bin/bash

# Job Flags
#SBATCH -p mit_preemptable
#SBATCH -c 4
#SBATCH --mem=512G
#SBATCH --requeue

# Set up environment
module load miniforge

conda activate james_linc_convert

# Run your application
linc-convert lsm stripes /orcd/data/linc/001/js1518/001412/derivatives/sub-MF283/strips_raw/sample-slice036 --out "/orcd/data/linc/001/js1518/sub-MF283_stripes/" --y_start 0.0 --y_end 800.0

Then I would generate vertical slices using a script that resembles

#!/bin/bash

#SBATCH -p mit_preemptable
#SBATCH -c 32
#SBATCH --mem=128G
#SBATCH --array=0-5%3
#SBATCH --requeue

# Set up environment
module load miniforge
conda activate james_linc_convert

# Parameters
CHUNK_SIZE=32
TASK_ID=$SLURM_ARRAY_TASK_ID

START=$((TASK_ID * CHUNK_SIZE))
END=$((START + CHUNK_SIZE))

OUTPUT_PATH="/orcd/data/linc/001/js1518/final_036_${TASK_ID}.ome.zarr"

echo "Task ID: $TASK_ID"
echo "x_chunk_start=$START, x_chunk_end=$END"
echo "Output: $OUTPUT_PATH"

# Run your application
linc-convert lsm stitch /orcd/data/linc/001/js1518/001412/derivatives/sub-MF283/strips_raw/sample-slice036/ \
  --out "$OUTPUT_PATH" \
  --levels 2 --nii --overlap 278 --blend \
  --y_start 0 --y_end 800 \
  --number_workers 4 --threads_per_worker 8 \
  --allow_padding --chunk 256 \
  --voxel_size "[1.0, 1.16, 1.16]" \
  --skew_angle 36 \
  --chunks_processed 32 \
  --x_chunk_start $START --x_chunk_end $END \
  --stripes /orcd/data/linc/001/js1518/sub-MF283_stripes/ \
  --checkpoint_file "/orcd/data/linc/001/js1518/final_036_checkpoint${TASK_ID}.dat" \
  --flip_vertically \
  --alternate_pattern

Then I would copy them all into one ome.zarr using a script like

#!/bin/bash

#SBATCH -p mit_preemptable
#SBATCH -c 4
#SBATCH --mem=32G

echo "Starting directory combine..."

INPUT_PARENT="/orcd/data/linc/001/js1518"
OUTPUT_ZARR="${INPUT_PARENT}/final_combined_036.ome.zarr"

mkdir -p "$OUTPUT_ZARR"

# Loop from 0_0 to 0_12
for i in {0..5}; do
    zarr="${INPUT_PARENT}/final_036_${i}.ome.zarr"

    echo "Processing $zarr"

    # Skip if directory doesn't exist
    [ -d "$zarr" ] || continue

    find "$zarr" -type f | while read file; do
        rel="${file#$zarr/}"
        dest="$OUTPUT_ZARR/$rel"

        mkdir -p "$(dirname "$dest")"

        # Fast move (no copy if same filesystem)
        mv "$file" "$dest"
    done

    # Clean up empty directories
    #find "$zarr" -type d -empty -delete
done

echo "Combine complete."

Then I would expand the pyramid with the following script

#!/bin/bash

#SBATCH -p mit_preemptable
#SBATCH -c 32
#SBATCH --mem=128G
#SBATCH --requeue

# Set up environment
module load miniforge
conda activate james_linc_convert

# Parameters
CHUNK_SIZE=32

START=$((TASK_ID * CHUNK_SIZE))
END=$((START + CHUNK_SIZE))

OUTPUT_PATH="/orcd/data/linc/001/js1518/final_combined_036.ome.zarr"

# Run your application
linc-convert lsm pyramid "$OUTPUT_PATH" \
  --start_level 3 --end_level 7 \
  --voxel_size "[1.0, 1.16, 1.16]" \
  --chunk 256