Skip to content

Commit

Permalink
Remove dead code
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 716480471
  • Loading branch information
tamaranorman authored and Torax team committed Jan 17, 2025
1 parent df71c61 commit 48c5eae
Show file tree
Hide file tree
Showing 11 changed files with 980 additions and 1,152 deletions.
18 changes: 18 additions & 0 deletions torax/config/runtime_params_slice.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from torax.config import profile_conditions
from torax.config import runtime_params as general_runtime_params_lib
from torax.geometry import geometry
from torax.geometry import geometry_provider as geometry_provider_lib
from torax.pedestal_model import runtime_params as pedestal_model_params
from torax.sources import runtime_params as sources_params
from torax.stepper import runtime_params as stepper_params
Expand Down Expand Up @@ -345,6 +346,23 @@ def __call__(
)


def get_consistent_dynamic_runtime_params_slice_and_geometry(
*,
t: chex.Numeric,
dynamic_runtime_params_slice_provider: DynamicRuntimeParamsSliceProvider,
geometry_provider: geometry_provider_lib.GeometryProvider,
) -> tuple[DynamicRuntimeParamsSlice, geometry.Geometry]:
"""Returns the dynamic runtime params and geometry for a given time."""
geo = geometry_provider(t)
dynamic_runtime_params_slice = dynamic_runtime_params_slice_provider(
t=t,
)
dynamic_runtime_params_slice, geo = make_ip_consistent(
dynamic_runtime_params_slice, geo
)
return dynamic_runtime_params_slice, geo


def make_ip_consistent(
dynamic_runtime_params_slice: DynamicRuntimeParamsSlice,
geo: geometry.Geometry,
Expand Down
12 changes: 6 additions & 6 deletions torax/config/tests/build_sim.py
Original file line number Diff line number Diff line change
Expand Up @@ -291,12 +291,12 @@ def test_chease_geometry_updates_Ip(self):
torax_mesh=geo_provider.torax_mesh,
)
)
geo = geo_provider(t=0)
dynamic_runtime_params_slice = runtime_params_provider(
t=0,
)
dynamic_slice, geo = runtime_params_slice.make_ip_consistent(
dynamic_runtime_params_slice, geo
dynamic_slice, geo = (
runtime_params_slice.get_consistent_dynamic_runtime_params_slice_and_geometry(
t=0,
dynamic_runtime_params_slice_provider=runtime_params_provider,
geometry_provider=geo_provider,
)
)
self.assertIsInstance(geo, geometry.StandardGeometry)
self.assertIsNotNone(dynamic_slice)
Expand Down
29 changes: 0 additions & 29 deletions torax/fvm/cell_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,28 +76,6 @@ class CellVariable:
# Can't make the above default values be jax zeros because that would be a
# call to jax before absl.app.run

def project(self, weights):
assert self.history is not None

def project(x):
return jnp.dot(weights, x)

def opt_project(x):
if x is None:
return None
return project(x)

return dataclasses.replace(
self,
value=project(self.value),
dr=self.dr[0],
left_face_constraint=opt_project(self.left_face_constraint),
left_face_grad_constraint=opt_project(self.left_face_grad_constraint),
right_face_constraint=opt_project(self.right_face_constraint),
right_face_grad_constraint=opt_project(self.right_face_grad_constraint),
history=None,
)

def __post_init__(self):
self.sanity_check()

Expand Down Expand Up @@ -242,13 +220,6 @@ def face_value(self) -> jax.Array:
)
return jnp.concatenate([left_face, inner, right_face], axis=-1)

def grad(self) -> jax.Array:
"""Returns the gradient of this variable wrt cell centers."""

self.assert_not_history()
face = self.face_value()
return jnp.diff(face) / self.dr

def history_elem(self) -> CellVariable:
"""Return a history entry version of this CellVariable."""
return dataclasses.replace(self, history=True)
Expand Down
Loading

0 comments on commit 48c5eae

Please sign in to comment.