distributed-regrid.yaml

apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
  name: distributed-regrid
  annotations:
    workflows.argoproj.io/description: >-
      Templates to regrid larger xarray Zarr Store distributed across years,
      without an external dask cluster, given an fsspec URL.
    workflows.argoproj.io/tags: fsspec,regrid,zarr,xarray
    workflows.argoproj.io/version: '>= 3.1.0'
  labels:
    component: utils
spec:
  entrypoint: main
  arguments:
    parameters:
      - name: in-zarr
        value: "gs://in-path.zarr"
      - name: regrid-method
        value: "bilinear"
      - name: domain-file
        value: "gs://support-c23ff1a3/domain.0p25x0p25.zarr"
  templates:

    - name: main
      inputs:
        parameters:
          - name: in-zarr
          - name: regrid-method
          - name: domain-file
          - name: add-cyclic-lon
            value: "false"
          - name: add-lat-buffer
            value: "false"
          # Space delimited list of all variables/coords in "in-zarr" without need for time dimension.
          - name: nontime-variables
            value: "lat lon"
      outputs:
        parameters:
          - name: out-zarr
            valueFrom:
              parameter: "{{ steps.prime-regrid-zarr.outputs.parameters.out-zarr }}"
      steps:
        - - name: prime-regrid-zarr
            template: prime-regrid-zarr
            arguments:
              parameters:
                - name: in-zarr
                  value: "{{ inputs.parameters.in-zarr }}"
                - name: regrid-method
                  value: "{{ inputs.parameters.regrid-method }}"
                - name: domain-file
                  value: "{{ inputs.parameters.domain-file }}"
                - name: nontime-variables
                  value: "{{ inputs.parameters.nontime-variables }}"
        - - name: regrid-year-to-primedzarr
            template: regrid-select-year-to-primedzarr
            arguments:
              parameters:
                - name: in-zarr
                  value: "{{ inputs.parameters.in-zarr }}"
                - name: out-zarr
                  value: "{{ steps.prime-regrid-zarr.outputs.parameters.out-zarr }}"
                - name: regrid-method
                  value: "{{ inputs.parameters.regrid-method }}"
                - name: domain-file
                  value: "{{ inputs.parameters.domain-file }}"
                - name: nontime-variables
                  value: "{{ inputs.parameters.nontime-variables }}"
                - name: select-year
                  value: "{{ item }}"
                - name: add-cyclic-lon
                  value: "{{ inputs.parameters.add-cyclic-lon }}"
                - name: add-lat-buffer
                  value: "{{ inputs.parameters.add-lat-buffer }}"
            withSequence:
              start: "{{ steps.prime-regrid-zarr.outputs.parameters.firstyear }}"
              end: "{{ steps.prime-regrid-zarr.outputs.parameters.lastyear }}"


    # Sets up output zarr file and metadata without writing time-related data to the zarr store.
    - name: prime-regrid-zarr
      inputs:
        parameters:
          - name: in-zarr
          - name: regrid-method
          - name: domain-file
          - name: nontime-variables
            value: "lat lon"
          - name: out-zarr
            value: "gs://scratch-170cd6ec/{{ workflow.uid }}/{{ pod.name }}/regridded.zarr"
      outputs:
        parameters:
          - name: out-zarr
            value: "{{ inputs.parameters.out-zarr }}"
          - name: firstyear
            valueFrom:
              path: "/tmp/firstyear.txt"
          - name: lastyear
            valueFrom:
              path: "/tmp/lastyear.txt"
      script:
        image: us-central1-docker.pkg.dev/downscalecmip6/private/dodola:0.8.0
        env:
          - name: IN_ZARR
            value: "{{ inputs.parameters.in-zarr }}"
          - name: OUT_ZARR
            value: "{{ inputs.parameters.out-zarr }}"
          - name: REGRID_METHOD
            value: "{{ inputs.parameters.regrid-method }}"
          - name: DOMAIN_FILE
            value: "{{ inputs.parameters.domain-file }}"
          - name: NONTIME_VARIABLES
            value: "{{ inputs.parameters.nontime-variables }}"
        command: [ python ]
        source: |
          import os
          import dodola.repository
          from dodola.core import xesmf_regrid
          import xarray as xr

          in_zarr = os.environ["IN_ZARR"]
          out_zarr = os.environ["OUT_ZARR"]
          regrid_method = os.environ["REGRID_METHOD"]
          domain_file = os.environ["DOMAIN_FILE"]
          # Space-delim string of variables and coordinates that do use "time" dimension.
          non_time = os.environ["NONTIME_VARIABLES"].strip().rsplit()

          ds_in = dodola.repository.read(in_zarr)
          domain_fl = dodola.repository.read(domain_file)

          first_year = ds_in["time"].dt.year.min().item()
          last_year = ds_in["time"].dt.year.max().item()
          print(f"{first_year=}, {last_year=}")

          ds_out = xesmf_regrid(ds_in, domain_fl, regrid_method, astype="float32")
          ds_out = ds_out.chunk({"time": 365, "lat": 2, "lon": -1})

          # Merge original attrs into output.
          ds_out.attrs |= ds_in.attrs
          for k, v in ds_in.variables.items():
              if k in ds_out:
                  ds_out[k].attrs |= v.attrs

          print(f"{ds_out=}")  # DEBUG

          # Output metadata to Zarr store.
          ds_out.to_zarr(
              out_zarr,
              mode="w",
              compute=False,
              consolidated=True
          )

          print(f"{ds_out[non_time]=}")  # DEBUG

          # Append variables that do not depend on "time"
          if non_time:
              ds_out[non_time].to_zarr(
                  out_zarr,
                  mode="a",
                  compute=True,
                  consolidated=True
              )

          with open("/tmp/firstyear.txt", mode="w") as fl:
              fl.write(str(first_year))
          with open("/tmp/lastyear.txt", mode="w") as fl:
              fl.write(str(last_year))
        resources:
          requests:
            memory: 4Gi
            cpu: "1000m"
          limits:
            memory: 4Gi
            cpu: "1000m"
      activeDeadlineSeconds: 1800
      retryStrategy:
        limit: 2
        retryPolicy: "Always"


    # Write select year of time-indexed data to a pre-primed zarr file.
    - name: regrid-select-year-to-primedzarr
      inputs:
        parameters:
          - name: in-zarr
          - name: select-year
          - name: regrid-method
          - name: domain-file
          - name: out-zarr
          - name: add-cyclic-lon
            value: "false"
          - name: add-lat-buffer
            value: "false"
          - name: nontime-variables
            value: "lat lon"
      outputs:
        parameters:
          - name: out-zarr
            value: "{{ inputs.parameters.out-zarr }}"
      script:
        image: us-central1-docker.pkg.dev/downscalecmip6/private/dodola:0.8.0
        env:
          - name: IN_ZARR
            value: "{{ inputs.parameters.in-zarr }}"
          - name: OUT_ZARR
            value: "{{ inputs.parameters.out-zarr }}"
          - name: SELTIME
            value: "{{ inputs.parameters.select-year }}"
          - name: REGRID_METHOD
            value: "{{ inputs.parameters.regrid-method }}"
          - name: DOMAIN_FILE
            value: "{{ inputs.parameters.domain-file }}"
          - name: NONTIME_VARIABLES
            value: "{{ inputs.parameters.nontime-variables }}"
          - name: ADD_CYCLIC_LON
            value: "{{ inputs.parameters.add-cyclic-lon }}"
          - name: ADD_LAT_BUFFER
            value: "{{ inputs.parameters.add-lat-buffer }}"
        command: [ python ]
        source: |
          import os
          import dodola.repository
          from dodola.core import xesmf_regrid
          import xarray as xr
          import numpy as np

          in_zarr = os.environ["IN_ZARR"]
          out_zarr = os.environ["OUT_ZARR"]
          sel_time = os.environ["SELTIME"]
          regrid_method = os.environ["REGRID_METHOD"]
          domain_file = os.environ["DOMAIN_FILE"]
          # Space-delim string of variables and coordinates that do use "time" dimension.
          # If data has something like `lat_b` or `height`, and you care about getting it in output,
          # those names need to be in this variable!
          # TODO: We could make this dynamically find non-time dependent variables/coord variables.
          non_time = os.environ["NONTIME_VARIABLES"].strip().rsplit()
          add_cyclic_lon =  os.environ["ADD_CYCLIC_LON"].lower() == "true"
          add_lat_buffer =  os.environ["ADD_LAT_BUFFER"].lower() == "true"

          # Setup regrid calculation for a time range.
          # Clip out target year from input data. This is a mess because we need the matches AND their index along
          # time, given the full dataset so that we can write back to the correct region in the output zarr store.
          ds_in = dodola.repository.read(in_zarr)

          ds_in = ds_in.sel(time=sel_time)
          ds_in.load()
          domain_fl = dodola.repository.read(domain_file)
          domain_fl.load()

          # Append mean of highest and lowest latitude band to data arrays
          # to cover north and south poles in global datasets. Assumes lat is
          # sorted from lowest to highest.
          if add_lat_buffer:
              # What are steps between latitudes. Rounding to avoid getting multiple
              # unique values just from float precision. Assumes regular steps and nothing fancy:
              delta = np.unique(ds_in["lat"].diff("lat").round(decimals=2).data).item()

              # Append additional latitude for pole that is mean of lowest latitude.
              min_lat = ds_in.isel(lat=0)
              min_lat = xr.ones_like(min_lat) * min_lat.mean(dim="lon")
              min_lat["lat"].data -= delta

              # Same but for other hemisphere pole
              max_lat = ds_in.isel(lat=-1)
              max_lat = xr.ones_like(max_lat) * max_lat.mean(dim="lon")
              max_lat["lat"].data += delta

              ds_in = xr.concat([min_lat, ds_in, max_lat], dim="lat")

          if add_cyclic_lon:
              ds_out = xesmf_regrid(ds_in, domain_fl, regrid_method, astype="float32", add_cyclic="lon")
          else:
              ds_out = xesmf_regrid(ds_in, domain_fl, regrid_method, astype="float32")

          ds_out = ds_out.chunk({"time": 365, "lat": 2, "lon": -1})

          with xr.open_zarr(out_zarr) as out_store:
              ds_out.attrs |= out_store.attrs
              for k, v in out_store.variables.items():
                  if k in ds_out:
                      ds_out[k].attrs |= v.attrs
              target_idx_slice = out_store["time"].to_index().get_loc(sel_time)

          if non_time:
              ds_out = ds_out.drop_vars(non_time)

          # Write to isolated region of Zarr store so can be done by independent processes.
          ds_out.to_zarr(out_zarr, region={"time": target_idx_slice}, mode="a")
        resources:
          requests:
            memory: 16Gi
            cpu: "1000m"
          limits:
            memory: 16Gi
            cpu: "4000m"
      activeDeadlineSeconds: 900
      retryStrategy:
        limit: 2
        retryPolicy: "Always"
        backoff:
          duration: 30s
          factor: 2