Code for TEAMS and other ensemble algorithms for Finkel et al 2025

This repository contains code for running the TEAMS algorithm, as well as more general ensemble-generating procedures for dynamical systems. Among the included examples are the Frierson GCM featured in Finkel & O'Gorman (2025), as well as the Lorenz-96 system featured in Finkel & O'Gorman (2024). The summary below is meant to convey the general design of the codebase, which aims to be extensible to new algorithms and new dynamical systems. But the implementation is not professional, and will certainly need case-by-case modification. Interested readers are highly encouraged to contact Justin Finkel (justinfocus12@gmail.com) for guidance on using it. Still, we hope that the core ideas conveyed here, in the publications, and in the code structure will help to facilitate community uptake of our methods.

• The driver code hierarchy: dynamical systems and algorithms for running them in ensemble mode
  • DynamicalSystem class (dynamicalsystem.py): specification of required methods for defining any new system, the most important one being run_trajectory. This method takes four arguments beyond the implicit "self" parameter (self gives the rule for stepping forward in time).
    • The first argument, "icandf", stands for "initial conditions and forcing", and contains all the needed information to (deterministically) provide the initial condition and generate noise, be it in the form of a one-off kick to the initial condition, a sequence of such perturbations, or an integer seed for generating a white noise sequence. All these versions can be found in the Forcing class (forcing.py), and one can implement more versions of forcing such as Levy noise if one wanted to. Because an icandf dictionary must be supplied to every run, any concrete implementation of DynamicalSystem must implement its own generate_default_icandf method. The role of this "icandf" dictionary is crucial, as it provides the extra input necessary beyond the evolution rule The same file implements "ODESystem" and "SDESystem" for ordinary and stochastic differential equations. An SDESystem instance has an ODESystem instance as an instance variable(e.g., the SDE $dX = -aX dt + b dW$, has the ODE $dx/dt = -ax$ as an instance variable). Both of those must be specified further for particular systems, such as Lorenz-96 as done in the examples. If ">" means "is a superclass of", DynamicalSystem > SDESystem > Lorenz96. For the Frierson GCM, more customization is necessary, and so the FriersonGCM class implemented in examples/frierson_gcm directly extends DynamicalSystem: DynamicalSystem > FriersonGCM.
    • The second argument is obs_fun: a function handle for any observable of special interest that we should evaluate online while running the model. Technically mandatory, but you could just have it return None.
    • The third argument is "saveinfo": a dictionary that specified how to save out the data, including possibly the initial condition and different fields in different output files. Its structure is different for each dynamical system.
    • The fourth argument is root_dir: the directory relative to which all output data is stored. This is helpful if one needs to migrate the entire ensemble from one location to another and only have to change one variable in the objects saved to disk.
  • Ensemble class (ensemble.py) specifies a data structure for managing large, and growing, ensembles of simulations stored on the disk. It uses a networkx graph to represent relationships between ensemble members, and is agnostic to the type of dynamical system. Any Ensemble instance has a DynamicalSystem instance as an attribute. Any call to DynamicalSystem.run_trajectory must go through the Ensemble instance that contains the DynamicalSystem instance, via the branch_or_plant method. Note there is a lot of dendromorphizing in the nomenclature (that's anthropomorphizing, but with tree-like attributes instead of human-like attributes).
  • EnsembleAlgorithm class (algorithms.py): in general, an EnsembleAlgorithm is a rule for generating the next member in an Ensemble given its current state. TEAMS is one example; DNS is another example with a very simple rule of extending the previous simulation chunk, picking it up at the endpoint; and "PeriodicBranching" is another example that we used to calculate ensemble dispersion rates in Finkel and O'Gorman (2025). Each of these is specified as a still-abstract subclass in the algorithms.py file, and the concrete implementations are completed in the examples subfolder. E.g., FriersonGCMTEAMS < TEAMS < EnsembleAlgorithm, and Lorenz96TEAMS < SDETEAMS < TEAMS < EnsembleAlgorithm.
• The Lorenz96 system, implemented in examples/lorenz96, is a good example to follow if implementing an algorithm on a new SDESystem.
• The Frierson GCM
  • As a DynamicalSystem object, it is implemented as FriersonGCM in examples/frierson_gcm/frierson_gcm.py. The core methods like run_trajectory merely wrap, via system calls, a large Fortran codebase that exists in some other directory, which is specified in a "config" dictionary, which must be supplied when creating an instance of FriersonGCM (see __init__ method). See the default_config method for a template. In our case, that codebase is in the companion repository jf_conv_gray_smooth, on the sppt branch. I will not attempt to describe that repository in detail --- it comes packaged with professional documentation from GFDL --- but point out that the main interface with this model happens in a few important methods:
    • compile_mppnccombine compiles the code for converting output to netcdf
    • generate_default_icandf demonstrates the two possible forcing choices: SPPT (which we used in Finkel & O'Gorman 2025), and impulsive (INT) forcing.
    • default_config: Here you can adjust things like the resolution (T21 or T42, which will automatically result in updates to the namelist)
    • derive_parameters: set some instance variables following the config dictionary
  • Within examples/frierson_gcm, the file algorithms_frierson.py implements concrete versions of abstract Algorithms classes, most relevantly FriersonGCMTEAMS < TEAMS, FriersonGCMPeriodicBranching < PeriodicBranching, and FriersonGCMDirectNumericalSimulation < DirectNumericalSimulation. For the former two, one of the inputs to __init__ are init_conds (restart files) and init_times (integer timestamps). There are various factory methods initialize_from... that take in a "dns" parameter which takes the form of a FriersonGCMDirectNumericalSimulation instance, and find appropriate initial conditions from that DNS.
    • DNS on the Frierson GCM: the script dns_frierson.py actually runs and analyzes DNS. The __main__ block looks for a command line argument which specifies the "procedure": either "single" (run the GCM with a single set of physical parameters), or "meta" (compare a collection of completed runs across different parameters). In the case of "single", a second argument is expected: an integer index i_expt specifying which parameter value to use. The function dns_multiparams creates the multidimensional array of parameters, and i_expt codes for a flat index (row-major / C-style ordering). In its present form, three of those parameters are singletons and only "sigmas" is a variable parameter; these are the noise strengths for SPPT. The other scripts (pebr_frierson.py and teams_frierson.py) have a similar setup, but with genuinely multidimensional parameter sets, e.g., advance split time and target field. Converting the flat index to a multi-index facilitates running these experiments in large batches on SLURM. A call looks like: $ python dns_frierson.py single 4 where the "4" specifies to use sigma=0.3.
    • TEAMS on the FriersonGCM: the script teams_frierson.py runs and analyzes TEAMS. The __main__ block has four procedures: single (a single run with a single seed), multiseed (aggregate analysis of many independent runs with the same parameters and different seeds, e.g., return period plots), and multidelta (analysis of performance across different advance split times). They should be executed in that order, and all after DNS so there are initial conditions to draw ancestors from. A call looks like: $ python teams_frierson.py single 281 (281 is arbitrary, but represents a flattened index in the space of target variable {precip, temperature} x advance split time {0, 4, 6, ..., 20} x seed {0, 1, ..., 47}.
    • PeriodicBranching on the Frierson GCM: the script pebr_frierson.py runs and analyzes PeriodicBranching. Here, the 'meta' procedure is outdated and needs updating. A call looks like $ python pebr_frierson.py single 4 The scripts described above share a common procedural structure, where <algorthmname>_single_procedure functions have a dictionary tododict of tasks assigned 0 or 1, so that you can run different stages in the pipeline in isolation. This is my way of implementing a notebook-like workflow in a script.