Skip to content

Framework API

The src.framework module defines the abstract Planet base class and all state dataclasses. Concrete planets (e.g. Mars) inherit from these and implement planet-specific physics.

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
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_ice]."""
        return torch.stack([
            self.thermal.surface_temperature,
            self.atmosphere.surface_pressure,
            self.water.ice_mass,
        ])

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

        Values are clamped to physical bounds.  Uses ``torch.where`` instead
        of a Python ``if`` so this method is device-agnostic and compatible
        with batched / compiled execution.
        """
        self.thermal.surface_temperature = y[0].clamp(min=1.0)
        self.atmosphere.surface_pressure = y[1].clamp(min=0.0)
        new_total = y[2].clamp(min=0.0)

        old_total = self.water.ice_mass_north + self.water.ice_mass_south
        # Apportion change between north/south reservoirs proportionally.
        # Use torch.where to avoid a Python boolean check on a GPU tensor.
        safe_total = old_total.clamp(min=1e-30)
        f_n = torch.where(
            old_total > 0.0,
            self.water.ice_mass_north / safe_total,
            old_total.new_full((), 0.5),  # same device & dtype, scalar shape
        )
        f_s = 1.0 - f_n
        self.water.ice_mass_north = (new_total * f_n).clamp(min=0.0)
        self.water.ice_mass_south = (new_total * f_s).clamp(min=0.0)
        self.water.ice_mass       = new_total

    # ------------------------------------------------------------------
    # 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_ice].

Source code in package/src/framework/planet.py 
120
121
122
123
124
125
126
def pack_state(self) -> torch.Tensor:
    """Pack the evolvable variables into a 1-D tensor [T, P, M_ice]."""
    return torch.stack([
        self.thermal.surface_temperature,
        self.atmosphere.surface_pressure,
        self.water.ice_mass,
    ])

unpack_state(y)

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

Values are clamped to physical bounds. Uses torch.where instead of a Python if so this method is device-agnostic and compatible with batched / compiled execution.

Source code in package/src/framework/planet.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def unpack_state(self, y: torch.Tensor) -> None:
    """Unpack a 1-D tensor [T, P, M_ice] back into the planet.

    Values are clamped to physical bounds.  Uses ``torch.where`` instead
    of a Python ``if`` so this method is device-agnostic and compatible
    with batched / compiled execution.
    """
    self.thermal.surface_temperature = y[0].clamp(min=1.0)
    self.atmosphere.surface_pressure = y[1].clamp(min=0.0)
    new_total = y[2].clamp(min=0.0)

    old_total = self.water.ice_mass_north + self.water.ice_mass_south
    # Apportion change between north/south reservoirs proportionally.
    # Use torch.where to avoid a Python boolean check on a GPU tensor.
    safe_total = old_total.clamp(min=1e-30)
    f_n = torch.where(
        old_total > 0.0,
        self.water.ice_mass_north / safe_total,
        old_total.new_full((), 0.5),  # same device & dtype, scalar shape
    )
    f_s = 1.0 - f_n
    self.water.ice_mass_north = (new_total * f_n).clamp(min=0.0)
    self.water.ice_mass_south = (new_total * f_s).clamp(min=0.0)
    self.water.ice_mass       = new_total

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 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
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 
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
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