Add a time dependent geometry provider

PiperOrigin-RevId: 653165315

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