File size: 5,099 Bytes
f4a39ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 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 | # Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PhysicsSpecifications contain physical parameters of dynamical systems.
To ensure that all model components the expected PhysicsSpecs all modules
(except specializing on a particular equation) must instantiate
PhysicsSpecs objects using `get_physics_specs`, which should be configured
appropriately via `gin`.
"""
from typing import Sequence, Union
from dinosaur import primitive_equations
from dinosaur import scales
from dinosaur import shallow_water
import gin
import numpy as np
# TODO(jamieas): consolidate with `PrimitiveEquationSpecs`. In particular,
# decide whether 'specs' should have units or be nondimensionalized.
QuantityOrStr = Union[str, scales.Quantity]
GET_DEFAULT_SCALE = gin.external_configurable(
lambda: scales.DEFAULT_SCALE, name='GET_DEFAULT_SCALE')
GET_ATMOSPHERIC_SCALE = gin.external_configurable(
lambda: scales.ATMOSPHERIC_SCALE, name='GET_ATMOSPHERIC_SCALE')
@gin.configurable
def get_physics_specs(construct_fn=gin.REQUIRED):
"""Returns physical parameters object generated by `construct_fn`."""
return construct_fn() # pyrefly: ignore[not-callable]
@gin.register
def shallow_water_specs_constructor(
density_vals: Union[Sequence[float], np.ndarray],
density_units: QuantityOrStr = scales.WATER_DENSITY,
radius_si: QuantityOrStr = scales.RADIUS,
angular_velocity_si: QuantityOrStr = scales.ANGULAR_VELOCITY,
gravity_acceleration_si: QuantityOrStr = scales.GRAVITY_ACCELERATION,
scale: scales.Scale = scales.DEFAULT_SCALE
) -> shallow_water.ShallowWaterSpecs:
"""Constructs `ShallowWaterSpecs` using gin-configurable parameters.
Args:
density_vals: density values for each layer of the shallow water system.
density_units: units in which `density_vals` are specified.
radius_si: radius of the domain specified with units attached.
angular_velocity_si: angular velocity of the domain with units attached.
gravity_acceleration_si: gravity on the surface with units attached.
scale: a scale object specifying the scales to use for nondimensionalizing.
Returns:
ShallowWaterSpecs object containing physical parameters of the system.
"""
densities = np.asarray(density_vals) * scales.Quantity(density_units)
return shallow_water.ShallowWaterSpecs.from_si(
densities=densities, # pyrefly: ignore[unexpected-keyword]
radius_si=scales.Quantity(radius_si),
angular_velocity_si=scales.Quantity(angular_velocity_si),
gravity_acceleration_si=scales.Quantity(gravity_acceleration_si),
scale=scale)
@gin.register
def primitive_eq_specs_constructor(
radius_si: QuantityOrStr = scales.RADIUS,
angular_velocity_si: QuantityOrStr = scales.ANGULAR_VELOCITY,
gravity_acceleration_si: QuantityOrStr = scales.GRAVITY_ACCELERATION,
ideal_gas_constant_si: QuantityOrStr = scales.IDEAL_GAS_CONSTANT,
water_vapor_gas_constant_si: QuantityOrStr = scales.IDEAL_GAS_CONSTANT_H20,
water_vapor_isobaric_heat_capacity_si: QuantityOrStr = (
scales.WATER_VAPOR_CP),
kappa_si: QuantityOrStr = scales.KAPPA,
scale: scales.Scale = scales.DEFAULT_SCALE,
) -> primitive_equations.PrimitiveEquationsSpecs:
"""Constructs `PrimitiveEquationsSpecs` using gin-configurable parameters.
Args:
radius_si: radius of the domain with units attached.
angular_velocity_si: angular velocity of the domain with units attached.
gravity_acceleration_si: gravity on the surface with units attached.
ideal_gas_constant_si: the gas constant with units attached.
water_vapor_gas_constant_si: the gas constant for vapor with units attached.
water_vapor_isobaric_heat_capacity_si: isobaric heat capacity of vapor with
units attached.
kappa_si: `ideal_gas_constant / Cp` where Cp is the isobaric heat capacity.
scale: a scale object specifying the scales to use for nondimensionalizing.
Returns:
PrimitiveEquationsSpecs object containing physical parameters of the system.
"""
return primitive_equations.PrimitiveEquationsSpecs.from_si(
radius_si=scales.Quantity(radius_si),
angular_velocity_si=scales.Quantity(angular_velocity_si),
gravity_acceleration_si=scales.Quantity(gravity_acceleration_si),
ideal_gas_constant_si=scales.Quantity(ideal_gas_constant_si),
water_vapor_gas_constant_si=scales.Quantity(water_vapor_gas_constant_si),
water_vapor_isobaric_heat_capacity_si=scales.Quantity(
water_vapor_isobaric_heat_capacity_si),
kappa_si=scales.Quantity(kappa_si),
scale=scale)
|