diff --git a/torax/geometry_provider.py b/torax/geometry_provider.py index 9f653e31..e51be5c4 100644 --- a/torax/geometry_provider.py +++ b/torax/geometry_provider.py @@ -81,10 +81,7 @@ class ConstantGeometryProvider(GeometryProvider): def __init__(self, geo: geometry.Geometry): self._geo = geo - def __call__( - self, - t: chex.Numeric, - ) -> geometry.Geometry: + def __call__(self, t: chex.Numeric) -> geometry.Geometry: # The API includes time as an arg even though it is unused in order # to match the API of a GeometryProvider. del t # Ignored. @@ -93,3 +90,17 @@ def __call__( @property def torax_mesh(self) -> geometry.Grid1D: return self._geo.torax_mesh + + +class TimeDependentGeometryProvider(GeometryProvider): + """Returns a Geometry that changes over time.""" + + def __init__(self, geometry_provider: geometry.GeometryProvider): + self._geometry_provider = geometry_provider + + def __call__(self, t: chex.Numeric) -> geometry.Geometry: + return self._geometry_provider.get_geometry(t) + + @property + def torax_mesh(self) -> geometry.Grid1D: + return self._geometry_provider.torax_mesh