Pilot GSM functions to also operate on arrays

[aufdenkampe]

All modeling software that use these modules will need them computed over arrays of values.

The decision to first get these modules working correctly with scalars (i.e. a single value) was wise, because Fortran and C++ have lots of additional code to deal with arrays that could have substantially complicated refactoring.

Fortunately, scientific Python makes it very simple to "vectorize" calculations from single value calculations to calculations over arrays of values, which speeds up calculations 8000 times with very little change in coding. The standard for “vectorized” operations in Python uses the NumPy Library, but many Python libraries provide such functionality (including Xarray and Dask), and a Python Array API Standard in being developed.

As a pilot for "vectorizing" this package, let's start by modifying functions in GSM to accept and operate on either a single value float or an array of floats.

[sjordan29]

Baseline, we'll want to accept either floats or arrays for the general constituent concentration as input to this function. I feel like we would also want users to input floats or arrays for other inputs too (e.g., you can have one water temperature for the whole domain, or allow for variability and accept an array of values as well). However, I would also think there are other inputs that we would want to remain constant over the entire domain (e.g., you either want to turn settling on or off for the whole domain, but you wouldn't want settling allowed in some cells but not others).

Here's how I currently have it set up. @aufdenkampe - I'd appreciate your review of the type hints below. Those that are Union[float, np.array] I envision being allowed to vary spatially over a model domain. Those that only accept a single value should be constant over the entire model domain.

def GeneralConstituentKinetics(
                               GC: Union[float, np.array], 
                               TwaterC: Union[float, np.array], 
                               order: int, 
                               k_rc20: Optional[Union[float, np.array]] = 1.0, 
                               k_theta: Optional[Union[float, np.array]] = 1.047,
                               rgc_rc20: Optional[Union[float, np.array]] = 1.0,
                               rgc_theta: Optional[Union[float, np.array]] = 1.047,
                               release: Optional[bool] = True,
                               settling: Optional[bool] = False, 
                               settling_rate: Optional[Union[float, np.array]] = 0.1,
                               depth: Optional[Union[float, np.array]] = 1.0,
                               ) -> Union[float, np.array]:

    '''
    Compute a single general constituent

    Args:
        GC (float or np.array): General constituent concentration of a single cell or an array of cells
        TwaterC (float or np.array): Water temperature in degrees Celsius
        order (int): Order of reaction kinetics (0, 1, or 2)
        k_rc20 (float or np.array, optional): Reaction (decay) rate at 20 degrees Celsius
        k_theta (float or np.array, optional): Arrhenius temperature correction factor for decay rate
        rgc_rc20 (float or np.array, optional): Sediment release rate at 20 degrees Celsius
        rgc_theta (float or np.array, optional): Arrhenius temperature correction factor for sediment release
        release (bool, optional): Compute resuspension, True = on; False = off.
        settling (bool, optional): Compute setting (i.e., bed loss, on/off)
        settling_rate (float or np.array, optional): Settling rate (m/s)
        depth (float or np.array, optional): Depth of the bed (meters)

    Returns:
        dGCdt (float or np.array): Rate of change of general constituent concentration
    '''

[imscw95]

@sjordan29 Just a heads up, we need to change the formulation of this module to match the fortran's correctly. I have attached an screen shot from an email I sent about the inaccuracy in the python. The fortran is correct. I plan to work on this soon, but am busy with some other tasks at the moment. Will it affect your ability to vectorize the code?

[sjordan29]

Feedback from @PeterKlaver on comment:

• k_theta: global
• rgc_theta: global

[sjordan29]

@imscw95 - let's keep in touch on timelines. I'm also working on other tasks right now, so let's just chat right before you start to make sure we won't create any conflicts. I'm going to open up another issue so you can track your progress correcting this issue separately!

[xaviernogueira]

All computations will work on arrays thanks to the xarray powered refactor in PR #18