Update lambda function to write directly to zarr

[maxrjones]

This PR builds off #1 in two ways:

1. Move core functionality to a "magg" package to make changes simpler to test
2. Create a template Zarr store for holding output and write one chunk per variable from each worker

FYI a lot of the diff relates to running a linter and isn't meaningful changes

TODO:

• Limit template to Antarctic region
• Write to Icechunk rather than Zarr for improved reliability (e.g., if someone's trying to read data while updates are happening)

[maxrjones]

@espg do you know of a simple way to get the healpix and/or morton cell ids for all cells south of a latitude for constructing the template? Currently I build a template for the entire world which is wasteful.

[espg]

Yes -- have a look here.

There's two schemes for healpix, NESTED (which is what morton indexing uses), and RING, which is the default. The RING ordering is a constant enumeration of ID's, where the cell centers lie along a fixed latitude band; once the latitude band is exhausted, it increments to the next lower latitude band, and then wraps again.

Easiest thing to do is use the regular healpy libraries to get the RING coordinates, and then convert them to NESTED and/or morton indices. Since the RING scheme is just a linear increment, find the minimum index for a given latitude at the resolution, and then your can just slice to end of the coordinate ID for that resolution (since the last four cells will be those that touch the south pole).

[espg]

Maybe this and this is what you want?

That's casting to/from morton from/to an encoding called UNIQ ; that starts at the lowest resolution, and then just increments until the end of resolution, and then starts for the next level down. Healpy has tools to / from RING and NESTED.

So... you'd go latitude ⟶ RING ⟶ NESTED ⟶ UNIQ ⟶ morton ; last operation in mortie, first in healpy. Middle doesn't have a convenience function, but is straightforward:

# uniq = 4 * nside**2 + ipix

# Using healpy

import healpy as hp

# RING to UNIQ (need to go via NESTED)
def ring_to_uniq(nside, ipix_ring):
    ipix_nest = hp.ring2nest(nside, ipix_ring)
    return 4 * nside**2 + ipix_nest

# Note: UNIQ is defined for NESTED ordering, not RING. So you must convert RING → NESTED first.

# Using cdshealpix (more direct)

# from cdshealpix import healpix_to_uniq
import numpy as np

depth = 8  # order = log2(nside)
ipix_nest = np.array([0, 1, 2, 3])
uniq = healpix_to_uniq(depth, ipix_nest)

# Reverse (UNIQ → order + pixel)

def uniq_to_nest(uniq):
    order = (uniq.bit_length() - 3) // 2
    nside = 1 << order
    ipix_nest = uniq - 4 * nside**2
    return order, nside, ipix_nest   

Right now I just have functions to go to/from UNIQ... but all this is making me realize that I should probably have a few more functions to convert to/from NESTED and RING directly to morton. To at least make that conversion pipeline shorter.

[maxrjones]

Maybe this and this is what you want?

That's casting to/from morton from/to an encoding called UNIQ ; that starts at the lowest resolution, and then just increments until the end of resolution, and then starts for the next level down. Healpy has tools to / from RING and NESTED.

So... you'd go latitude ⟶ RING ⟶ NESTED ⟶ UNIQ ⟶ morton ; last operation in mortie, first in healpy. Middle doesn't have a convenience function, but is straightforward:

# uniq = 4 * nside**2 + ipix

# Using healpy

import healpy as hp

# RING to UNIQ (need to go via NESTED)
def ring_to_uniq(nside, ipix_ring):
    ipix_nest = hp.ring2nest(nside, ipix_ring)
    return 4 * nside**2 + ipix_nest

# Note: UNIQ is defined for NESTED ordering, not RING. So you must convert RING → NESTED first.

# Using cdshealpix (more direct)

# from cdshealpix import healpix_to_uniq
import numpy as np

depth = 8  # order = log2(nside)
ipix_nest = np.array([0, 1, 2, 3])
uniq = healpix_to_uniq(depth, ipix_nest)

# Reverse (UNIQ → order + pixel)

def uniq_to_nest(uniq):
    order = (uniq.bit_length() - 3) // 2
    nside = 1 << order
    ipix_nest = uniq - 4 * nside**2
    return order, nside, ipix_nest   

Right now I just have functions to go to/from UNIQ... but all this is making me realize that I should probably have a few more functions to convert to/from NESTED and RING directly to morton. To at least make that conversion pipeline shorter.

thanks, this should work well