Basics

Earth2Studio.jl holds the earth2studio Python module and its submodules as Py objects. Anything you can do in Python you can do in Julia by calling into these objects via PythonCall.jl.

Module references

Nothing is exported (upstream names such as run and io would clash with Base), so every reference is accessed through the module:

using Earth2Studio

Earth2Studio.earth2studio  # earth2studio
Earth2Studio.data          # earth2studio.data
Earth2Studio.models        # earth2studio.models  (use models.px and models.dx)
Earth2Studio.perturbation  # earth2studio.perturbation
Earth2Studio.io            # earth2studio.io
Earth2Studio.run           # earth2studio.run
Earth2Studio.statistics    # earth2studio.statistics
Earth2Studio.utils         # earth2studio.utils

Each is a PythonCall.Py reference — attribute access, calls, and conversions work the way they do in PythonCall. For shorter code, alias the module:

import Earth2Studio as E2S
E2S.data.ARCO()

Auto-conversion of common Julia types

PythonCall handles the conversions you most often need at call sites:

Julia	Python
DateTime	datetime.datetime
Vector{DateTime}	juliacall.VectorValue yielding datetime.datetime
String	str
Vector{String}	juliacall.VectorValue yielding str

A juliacall.VectorValue is not a Python list, but it supports len, indexing, and iteration. Each of those calls back into Julia. That is safe on the main thread, which is where the run workflows convert their time argument, so Julia vectors can be passed to them directly:

Earth2Studio.run.deterministic([DateTime(2023, 6, 15), DateTime(2023, 6, 16)], 10, model, ds, backend)

Data sources are different: they iterate time and variable inside an async fetch that runs on fsspec’s background IO thread. A callback into Julia from a thread Julia does not own can deadlock the process (Julia’s garbage collector waits for the main thread to reach a safepoint while the main thread waits for Python’s GIL). Convert with pylist before calling a data source:

using Earth2Studio
using Dates
using PythonCall: pylist

ds = Earth2Studio.data.ARCO()
da = ds(DateTime(2023, 6, 15), pylist(["t2m", "u10m"]))                # one time
da = ds(pylist([DateTime(2023, 6, 15), DateTime(2023, 6, 16)]), "t2m")  # several

Two practical caveats:

• Date (no time component) is converted to Python’s datetime.date, which earth2studio mixes with datetime.datetime internally and triggers TypeError. Wrap with DateTime(d) first.
• If a Python API specifically wants a numpy.datetime64 array, build it explicitly with pyimport("numpy").array(...).

Fetching data

using Earth2Studio
using PythonCall: pyconvert, pylist, PyArray
using Dates

ds = Earth2Studio.data.ARCO()
da = ds(DateTime(2023, 6, 15), pylist(["t2m", "u10m"]))

pyconvert(Tuple, da.dims)        # ("time", "variable", "lat", "lon")
pyconvert(Tuple, da.shape)       # (1, 2, 721, 1440)

arr = PyArray(da.values)         # zero-copy PyArray{Float64, 4}
t2m = arr[1, 1, :, :]            # 721 × 1440 lat × lon Kelvin

Variable names

earth2studio uses one set of short variable names across all data sources: surface fields such as t2m, u10m, msl, tp, and pressure-level fields written as <variable><level>, e.g. t850, z500, q700. Each data source has a lexicon class in earth2studio.lexicon whose VOCAB dict maps these names to the source’s native names, so its keys are the variables that source can provide. See the upstream lexicon guide for the full naming convention.

using Earth2Studio
using PythonCall: pyconvert, pycontains

lex = Earth2Studio.earth2studio.lexicon

vocab = lex.ARCOLexicon.VOCAB
length(vocab)                                    # 1436
pyconvert(Vector{String}, vocab.keys())[1:5]     # ["u10m", "v10m", "u100m", "v100m", "t2m"]
pycontains(vocab, "t850")                        # true

vocab["z500"]                                    # "geopotential::500"
lex.GFSLexicon.VOCAB["z500"]                     # "HGT::500 mb"

Prognostic models

pkg   = Earth2Studio.models.px.FCN.load_default_package()
model = Earth2Studio.models.px.FCN.load_model(pkg)
model.to("cuda")

Diagnostic models

pkg = Earth2Studio.models.dx.PrecipitationAFNO.load_default_package()
dx  = Earth2Studio.models.dx.PrecipitationAFNO.load_model(pkg)

Workflows

using Dates

backend = Earth2Studio.io.ZarrBackend("/tmp/forecast.zarr")

# Deterministic
Earth2Studio.run.deterministic([DateTime(2023, 6, 15)], 10, model, ds, backend)

# Ensemble
perturb = Earth2Studio.perturbation.SphericalGaussian()
Earth2Studio.run.ensemble([DateTime(2023, 6, 15)], 10, 4, model, ds, backend, perturb; batch_size=4)

# Diagnostic
Earth2Studio.run.diagnostic([DateTime(2023, 6, 15)], 10, model, dx, ds, backend)

IO backends

Earth2Studio.io.ZarrBackend("/tmp/out.zarr")
Earth2Studio.io.NetCDF4Backend("/tmp/out.nc")
Earth2Studio.io.AsyncZarrBackend("/tmp/out.zarr")
Earth2Studio.io.XarrayBackend()
Earth2Studio.io.KVBackend()

Statistics

m = Earth2Studio.statistics.rmse(reduction_dimensions=["time"])
m = Earth2Studio.statistics.acc(reduction_dimensions=["time"])
m = Earth2Studio.statistics.crps(ensemble_dimension="ensemble")

Bringing results back to Julia

using PythonCall

# numpy / xarray array → Julia
a_jl   = pyconvert(Array, da.values)   # copies
a_view = PyArray(da.values)            # zero-copy view

To keep dimension names and coordinates, convert to a DimArray instead; see the Julia packages tutorial.

Refer to the PythonCall.jl docs for more conversion utilities, and to the earth2studio docs for the Python API itself.