Skip to content

Framework API

The primary GCM uses src.framework.gcm for body constants, grids, dynamics, restart and parameterized integration; src.framework.physics for column operators; and src.framework.neural for learned components. Begin with the GCM API and architecture.

Deprecated global-mean framework

Deprecated torch model

The Planet base class and subsystem dataclasses below support the deprecated global-mean workflow. They remain available for existing users and are not the state/integration contract of the GCM.

Planet (Abstract Base Class)

src.framework.planet

Abstract Planet framework — PyTorch backend.

Defines the abstract base class Planet, the PlanetaryState snapshot, and OrbitalParameters. Concrete planets (e.g. Mars) inherit from Planet and provide planet-specific physics (derivatives, equilibrium calculations), while the integration strategy (RK4, reduced-order) lives in the engine (TimeController).

Separation of concerns

Planet → what the physics equations are (state + derivatives) Engine → how those equations are integrated (RK4 / relaxation)

All numerical values are stored as torch.Tensor scalars (dtype=torch.float64).

Planet

Bases: ABC

Abstract base class for planetary bodies.

A Planet is a property container combined with a physics model. It knows what the governing equations are (derivatives, equilibrium) but does not know how to integrate them — that is the engine's job.

Sub-classes must implement: setup_properties() – initialize the planet's property dataclasses compute_derivatives(y) – ODE RHS for the coupled system compute_fast_physics(dt) – reduced-order analytic update

Source code in package/src/framework/planet.py 
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
class Planet(ABC):
    """Abstract base class for planetary bodies.

    A ``Planet`` is a **property container** combined with a **physics model**.
    It knows *what* the governing equations are (derivatives, equilibrium)
    but does **not** know *how* to integrate them — that is the engine's job.

    Sub-classes **must** implement:
        ``setup_properties()``       – initialize the planet's property dataclasses
        ``compute_derivatives(y)``   – ODE RHS for the coupled system
        ``compute_fast_physics(dt)`` – reduced-order analytic update
    """

    # --- Concrete attributes set by subclass __init__ ---
    intrinsic_params: IntrinsicParameters
    orbital_params: OrbitalParameters
    _device: torch.device          # set by subclass before setup_properties()

    # --- Planetary Properties ---
    atmosphere: Atmosphere
    thermal: Thermal
    water: Water
    radiation: Radiation
    magnetic: Magnetic

    # --- Time tracking ---
    elapsed_time: torch.Tensor
    orbital_angle: torch.Tensor

    # ------------------------------------------------------------------
    # Abstract interface  —  physics provided by each planet
    # ------------------------------------------------------------------
    @abstractmethod
    def setup_properties(self) -> None:
        """Initialize the planet's physical properties."""
        ...

    @abstractmethod
    def compute_derivatives(self, y: torch.Tensor) -> torch.Tensor:
        """Compute dy/dt for the coupled ODE system.

        Parameters
        ----------
        y : torch.Tensor, shape [3]
            State vector ``[T_surface, P_surface, M_ice]``.

        Returns
        -------
        torch.Tensor, shape [3]
            Time derivatives ``[dT/dt, dP/dt, dM_ice/dt]``.

        The engine calls this with intermediate state vectors during
        RK4 sub-steps.  Implementations should use ``self`` only
        for quantities that are *not* part of *y* (e.g. albedo,
        greenhouse factor, solar flux).
        """
        ...

    @abstractmethod
    def compute_fast_physics(self, dt: torch.Tensor) -> None:
        """Apply a single reduced-order physics step to the planet.

        This method computes analytic / linearised approximations
        (equilibrium temperature, exponential relaxation, first-order
        Euler updates for pressure and ice) and writes the results
        directly into ``self.state``.

        The engine is responsible for calling ``advance_orbit`` before
        this method, so ``self.state.solar_flux`` is already up-to-date.

        Parameters
        ----------
        dt : torch.Tensor
            Timestep in seconds.
        """
        ...

    # ------------------------------------------------------------------
    # State packing / unpacking  (used by the engine's ODE integrators)
    # ------------------------------------------------------------------
    def pack_state(self) -> torch.Tensor:
        """Pack the evolvable variables into a 1-D tensor [T, P, M_N, M_S].

        The two polar caps are carried as **independent** state variables so the
        ODE integrator (RK4) evolves each reservoir on its own. Collapsing them
        to a single total and re-apportioning afterwards destroys the seasonal
        anti-phase exchange (one cap grows while the other shrinks) and lets the
        per-cap sublimation gate mis-fire — which drained the caps to zero on the
        ACCURATE path.
        """
        return torch.stack([
            self.thermal.surface_temperature,
            self.atmosphere.surface_pressure,
            self.water.ice_mass_north,
            self.water.ice_mass_south,
        ])

    def unpack_state(self, y: torch.Tensor) -> None:
        """Unpack a 1-D tensor [T, P, M_N, M_S] back into the planet.

        Each cap is set directly from its own state component (clamped ≥ 0); the
        total ice mass is their sum. No proportional re-split — the reservoirs
        are prognostic and independent.
        """
        self.thermal.surface_temperature = y[0].clamp(min=1.0)
        self.atmosphere.surface_pressure = y[1].clamp(min=0.0)
        self.water.ice_mass_north = y[2].clamp(min=0.0)
        self.water.ice_mass_south = y[3].clamp(min=0.0)
        self.water.ice_mass = self.water.ice_mass_north + self.water.ice_mass_south

    def observed_surface_pressure(self) -> torch.Tensor:
        """Local (observed) surface pressure.

        Base planets have no local pressure phenomena, so this is the
        prognostic (mass-budget) pressure unchanged.  Subclasses may add
        diagnostic overlays (e.g. Mars's thermal tide) that redistribute
        mass locally without creating or destroying it globally.
        """
        return self.atmosphere.surface_pressure

    # ------------------------------------------------------------------
    # Shared helpers available to every planet
    # ------------------------------------------------------------------
    def solar_flux_at_distance(self, distance: torch.Tensor) -> torch.Tensor:
        """Solar irradiance at heliocentric *distance* (metres).

        Uses Python-float constants so the result lives on whatever device
        ``distance`` is on — no explicit ``.to(device)`` required.

        Equation (inverse-square law):
            F(d) = F₀ × (1 AU / d)²

        Reference:
        https://en.wikipedia.org/wiki/Solar_irradiance
        https://en.wikipedia.org/wiki/Solar_constant
        """
        return _SOLAR_CONST_1AU * (_AU_METRES / distance) ** 2

    def advance_orbit(self, dt: torch.Tensor) -> None:
        """Advance orbital angle (mean-motion approximation) and update
        solar flux.

        Called by the engine at each timestep *before* physics updates.
        Uses Python-float math constants so all arithmetic stays on whichever
        device the planet's tensors live on.

        Equations:
            dθ/dt = 2π / T_orbital          (mean motion)
            r(θ)  = a(1−e²)/(1+e cos θ)    (Kepler ellipse)
            F     = F₀ (1 AU / r)²          (inverse-square)
        """
        self.elapsed_time = self.elapsed_time + dt

        """The approximation here: this treats orbital_angle as the true anomaly (actual angular position on the ellipse)
        but advances it at the constant rate of the mean anomaly. In reality, the planet moves faster near perihelion
        and slower near aphelion (Kepler's second law — equal areas in equal times). For Mars with e = 0.0934,
        this introduces a phase error of up to ~±10° in solar longitude, shifting the timing of seasonal peaks by roughly 5–10 sols.
        """
        self.orbital_angle = torch.remainder(
            self.orbital_angle
            + _TWO_PI * dt / self.orbital_params.orbital_period,
            _TWO_PI,
        )

        distance = self.orbital_params.distance_from_sun(self.orbital_angle)
        self.radiation.solar_flux = self.solar_flux_at_distance(distance)

setup_properties() abstractmethod

Initialize the planet's physical properties.

Source code in package/src/framework/planet.py 
72
73
74
75
@abstractmethod
def setup_properties(self) -> None:
    """Initialize the planet's physical properties."""
    ...

compute_derivatives(y) abstractmethod

Compute dy/dt for the coupled ODE system.

Parameters

y : torch.Tensor, shape [3] State vector [T_surface, P_surface, M_ice].

Returns

torch.Tensor, shape [3] Time derivatives [dT/dt, dP/dt, dM_ice/dt].

The engine calls this with intermediate state vectors during RK4 sub-steps. Implementations should use self only for quantities that are not part of y (e.g. albedo, greenhouse factor, solar flux).

Source code in package/src/framework/planet.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
@abstractmethod
def compute_derivatives(self, y: torch.Tensor) -> torch.Tensor:
    """Compute dy/dt for the coupled ODE system.

    Parameters
    ----------
    y : torch.Tensor, shape [3]
        State vector ``[T_surface, P_surface, M_ice]``.

    Returns
    -------
    torch.Tensor, shape [3]
        Time derivatives ``[dT/dt, dP/dt, dM_ice/dt]``.

    The engine calls this with intermediate state vectors during
    RK4 sub-steps.  Implementations should use ``self`` only
    for quantities that are *not* part of *y* (e.g. albedo,
    greenhouse factor, solar flux).
    """
    ...

compute_fast_physics(dt) abstractmethod

Apply a single reduced-order physics step to the planet.

This method computes analytic / linearised approximations (equilibrium temperature, exponential relaxation, first-order Euler updates for pressure and ice) and writes the results directly into self.state.

The engine is responsible for calling advance_orbit before this method, so self.state.solar_flux is already up-to-date.

Parameters

dt : torch.Tensor Timestep in seconds.

Source code in package/src/framework/planet.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
@abstractmethod
def compute_fast_physics(self, dt: torch.Tensor) -> None:
    """Apply a single reduced-order physics step to the planet.

    This method computes analytic / linearised approximations
    (equilibrium temperature, exponential relaxation, first-order
    Euler updates for pressure and ice) and writes the results
    directly into ``self.state``.

    The engine is responsible for calling ``advance_orbit`` before
    this method, so ``self.state.solar_flux`` is already up-to-date.

    Parameters
    ----------
    dt : torch.Tensor
        Timestep in seconds.
    """
    ...

pack_state()

Pack the evolvable variables into a 1-D tensor [T, P, M_N, M_S].

The two polar caps are carried as independent state variables so the ODE integrator (RK4) evolves each reservoir on its own. Collapsing them to a single total and re-apportioning afterwards destroys the seasonal anti-phase exchange (one cap grows while the other shrinks) and lets the per-cap sublimation gate mis-fire — which drained the caps to zero on the ACCURATE path.

Source code in package/src/framework/planet.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
def pack_state(self) -> torch.Tensor:
    """Pack the evolvable variables into a 1-D tensor [T, P, M_N, M_S].

    The two polar caps are carried as **independent** state variables so the
    ODE integrator (RK4) evolves each reservoir on its own. Collapsing them
    to a single total and re-apportioning afterwards destroys the seasonal
    anti-phase exchange (one cap grows while the other shrinks) and lets the
    per-cap sublimation gate mis-fire — which drained the caps to zero on the
    ACCURATE path.
    """
    return torch.stack([
        self.thermal.surface_temperature,
        self.atmosphere.surface_pressure,
        self.water.ice_mass_north,
        self.water.ice_mass_south,
    ])

unpack_state(y)

Unpack a 1-D tensor [T, P, M_N, M_S] back into the planet.

Each cap is set directly from its own state component (clamped ≥ 0); the total ice mass is their sum. No proportional re-split — the reservoirs are prognostic and independent.

Source code in package/src/framework/planet.py 
137
138
139
140
141
142
143
144
145
146
147
148
def unpack_state(self, y: torch.Tensor) -> None:
    """Unpack a 1-D tensor [T, P, M_N, M_S] back into the planet.

    Each cap is set directly from its own state component (clamped ≥ 0); the
    total ice mass is their sum. No proportional re-split — the reservoirs
    are prognostic and independent.
    """
    self.thermal.surface_temperature = y[0].clamp(min=1.0)
    self.atmosphere.surface_pressure = y[1].clamp(min=0.0)
    self.water.ice_mass_north = y[2].clamp(min=0.0)
    self.water.ice_mass_south = y[3].clamp(min=0.0)
    self.water.ice_mass = self.water.ice_mass_north + self.water.ice_mass_south

observed_surface_pressure()

Local (observed) surface pressure.

Base planets have no local pressure phenomena, so this is the prognostic (mass-budget) pressure unchanged. Subclasses may add diagnostic overlays (e.g. Mars's thermal tide) that redistribute mass locally without creating or destroying it globally.

Source code in package/src/framework/planet.py 
150
151
152
153
154
155
156
157
158
def observed_surface_pressure(self) -> torch.Tensor:
    """Local (observed) surface pressure.

    Base planets have no local pressure phenomena, so this is the
    prognostic (mass-budget) pressure unchanged.  Subclasses may add
    diagnostic overlays (e.g. Mars's thermal tide) that redistribute
    mass locally without creating or destroying it globally.
    """
    return self.atmosphere.surface_pressure

solar_flux_at_distance(distance)

Solar irradiance at heliocentric distance (metres).

Uses Python-float constants so the result lives on whatever device distance is on — no explicit .to(device) required.

Equation (inverse-square law): F(d) = F₀ × (1 AU / d)²

Reference: https://en.wikipedia.org/wiki/Solar_irradiance https://en.wikipedia.org/wiki/Solar_constant

Source code in package/src/framework/planet.py 
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def solar_flux_at_distance(self, distance: torch.Tensor) -> torch.Tensor:
    """Solar irradiance at heliocentric *distance* (metres).

    Uses Python-float constants so the result lives on whatever device
    ``distance`` is on — no explicit ``.to(device)`` required.

    Equation (inverse-square law):
        F(d) = F₀ × (1 AU / d)²

    Reference:
    https://en.wikipedia.org/wiki/Solar_irradiance
    https://en.wikipedia.org/wiki/Solar_constant
    """
    return _SOLAR_CONST_1AU * (_AU_METRES / distance) ** 2

advance_orbit(dt)

Advance orbital angle (mean-motion approximation) and update solar flux.

Called by the engine at each timestep before physics updates. Uses Python-float math constants so all arithmetic stays on whichever device the planet's tensors live on.

Equations

dθ/dt = 2π / T_orbital (mean motion) r(θ) = a(1−e²)/(1+e cos θ) (Kepler ellipse) F = F₀ (1 AU / r)² (inverse-square)

Source code in package/src/framework/planet.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def advance_orbit(self, dt: torch.Tensor) -> None:
    """Advance orbital angle (mean-motion approximation) and update
    solar flux.

    Called by the engine at each timestep *before* physics updates.
    Uses Python-float math constants so all arithmetic stays on whichever
    device the planet's tensors live on.

    Equations:
        dθ/dt = 2π / T_orbital          (mean motion)
        r(θ)  = a(1−e²)/(1+e cos θ)    (Kepler ellipse)
        F     = F₀ (1 AU / r)²          (inverse-square)
    """
    self.elapsed_time = self.elapsed_time + dt

    """The approximation here: this treats orbital_angle as the true anomaly (actual angular position on the ellipse)
    but advances it at the constant rate of the mean anomaly. In reality, the planet moves faster near perihelion
    and slower near aphelion (Kepler's second law — equal areas in equal times). For Mars with e = 0.0934,
    this introduces a phase error of up to ~±10° in solar longitude, shifting the timing of seasonal peaks by roughly 5–10 sols.
    """
    self.orbital_angle = torch.remainder(
        self.orbital_angle
        + _TWO_PI * dt / self.orbital_params.orbital_period,
        _TWO_PI,
    )

    distance = self.orbital_params.distance_from_sun(self.orbital_angle)
    self.radiation.solar_flux = self.solar_flux_at_distance(distance)

State Dataclasses

Atmosphere

src.framework.atmosphere

Thermal

src.framework.thermal

Water / Cryosphere

src.framework.water

Orbital Parameters

src.framework.orbital

OrbitalParameters dataclass

Keplerian orbital elements (assumed constant over short runs).

Source code in package/src/framework/orbital.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
@dataclass
class OrbitalParameters:
    """Keplerian orbital elements (assumed constant over short runs)."""

    semi_major_axis: torch.Tensor      # metres
    eccentricity: torch.Tensor         # dimensionless, 0 < e < 1
    orbital_period: torch.Tensor       # seconds
    axial_tilt: torch.Tensor           # radians

    def distance_from_sun(self, theta: torch.Tensor) -> torch.Tensor:
        """Heliocentric distance at true anomaly *theta* (radians).

        Equation (Kepler's first law – conic section):
            r(θ) = a(1 − e²) / (1 + e cos θ)

        Reference: https://en.wikipedia.org/wiki/Kepler_orbit
        """
        return (
            self.semi_major_axis
            * (1.0 - self.eccentricity ** 2)
            / (1.0 + self.eccentricity * torch.cos(theta))
        )

distance_from_sun(theta)

Heliocentric distance at true anomaly theta (radians).

Equation (Kepler's first law – conic section): r(θ) = a(1 − e²) / (1 + e cos θ)

Reference: https://en.wikipedia.org/wiki/Kepler_orbit

Source code in package/src/framework/orbital.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
def distance_from_sun(self, theta: torch.Tensor) -> torch.Tensor:
    """Heliocentric distance at true anomaly *theta* (radians).

    Equation (Kepler's first law – conic section):
        r(θ) = a(1 − e²) / (1 + e cos θ)

    Reference: https://en.wikipedia.org/wiki/Kepler_orbit
    """
    return (
        self.semi_major_axis
        * (1.0 - self.eccentricity ** 2)
        / (1.0 + self.eccentricity * torch.cos(theta))
    )

Intrinsic Parameters

src.framework.intrinsic

Radiation

src.framework.radiation

Magnetic

src.framework.magnetic