<commit_before>#!/usr/bin/env python

import numpy as np

from landlab import Component
from .funcs import get_flexure_parameter


_VALID_METHODS = set(['airy', 'flexure'])


def assert_method_is_valid(method):
    if method not in _VALID_METHODS:
        raise ValueError('%s: Invalid method name' % method)


class FlexureComponent(Component):
    """
    Landlab component that implements a 1 and 2D lithospheric flexure
    model.

    Examples
    --------
    >>> from landlab import RasterModelGrid
    >>> from landlab.components.flexure import FlexureComponent
    >>> grid = RasterModelGrid(5, 4, 1.e4)
    >>> flex = FlexureComponent(grid)
    >>> flex.name
    'Flexure'
    >>> sorted(flex.input_var_names)
    ['lithosphere__elevation', 'lithosphere__overlying_pressure', 'planet_surface_sediment__deposition_increment']
    >>> sorted(flex.output_var_names)
    ['lithosphere__elevation', 'lithosphere__elevation_increment']
    >>> for var in sorted(flex.units): flex.units[var]
    'm'
    'm'
    'Pa'
    'm'

    >>> flex.grid.number_of_node_rows
    5
    >>> flex.grid.number_of_node_columns
    4
    >>> flex.grid is grid
    True

    >>> np.all(grid.at_node['lithosphere__elevation'] == 0.)
    True

    >>> np.all(grid.at_node['lithosphere__elevation'] == 0.)
    True
    >>> np.all(grid.at_node['lithosphere__overlying_pressure'] == 0.)
    True
    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere__elevation_increment'] == 0.)
    True

    >>> load = grid.at_node['lithosphere__overlying_pressure']
    >>> load[4] = 1e9
    >>> dz = grid.at_node['lithosphere__elevation_increment']
    >>> np.all(dz == 0.)
    True

    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere__elevation_increment'] == 0.)
    False
    """

    ###############THIS IS THE STANDARD DECLARATIONS EXAMPLE.
    #YOU NEED TO DECLARE ALL THE FOLLOWING LIKE THIS HERE IN YOUR COMPONENT!!
    _name = 'Flexure'
    #...give the component a name

    _input_var_names = set([
        'lithosphere__overlying_pressure',
        'lithosphere__elevation',
        'planet_surface_sediment__deposition_increment',
    ])
    #...the component requires these values to do its calculation

    _output_var_names = set([
        'lithosphere__elevation_increment',
        'lithosphere__elevation',
    ])
    #...the component modifies these values

    # the next three dictionaries must each have the union of the two var_name
    # sets above as their keys
    _var_units = {
        'lithosphere__overlying_pressure': 'Pa',
        'lithosphere__elevation': 'm',
        'lithosphere__elevation_increment': 'm',
        'planet_surface_sediment__deposition_increment': 'm',
    }
    #...the units for each field. In the future, there may be unit casting
    # required, as there's nothing stopping you passing in the wrong units!!!

    _var_mapping = {
        'lithosphere__overlying_pressure': 'node',
        'lithosphere__elevation': 'node',
        'lithosphere__elevation_increment': 'node',
        'planet_surface_sediment__deposition_increment': 'node',
    }
    #...the grid centering of each name. Note you CANNOT and SHOULD NOT TRY
    # TO define the same field name but with different centering within one
    # component; you'll need to distinguish the names.

    _var_doc = {
        'lithosphere__overlying_pressure': 'The pressure at the base of the lithosphere',
        'lithosphere__elevation': 'The elevation of the top of the lithosphere, i.e., the land surface',
        'lithosphere__elevation_increment': 'The change in elevation of the top of the lithosphere (the land surface) in one timestep',
        'planet_surface_sediment__deposition_increment': 'The amount of sediment deposited at the land surface in one timestep',
    }
    #...give the names a short description. [We still need to work out how
    # we deal with same name, but different description between components]

    ###lastly, ---> 1. did you remember to import and inherit from Component?
    #do the import up top ->
    #from landlab import Component
    #& in the class declaration ->
    #   class FlexureComponent(Component):
    #        ...
    ### ---> 2. Make sure that self._grid is an alias for the grid after
    # initialization, below
    ################################

    def __init__(self, grid, **kwds):
        self._eet = kwds.pop('eet', 65000.)
        self._youngs = kwds.pop('youngs', 7e10)
        self._method = kwds.pop('method', 'airy')
        self._grid = grid

        assert_method_is_valid(self._method)

        super(FlexureComponent, self).__init__(grid, **kwds)

        for name in self._input_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros('node', name, units=self._var_units[name])

        for name in self._output_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros('node', name, units=self._var_units[name])

        self._last_load = self.grid.field_values('node', 'lithosphere__overlying_pressure').copy()

        self._nodal_values = self.grid['node']

        self._r = self._set_kei_func_grid()

    def _set_kei_func_grid(self):
        from scipy.special import kei

        alpha = get_flexure_parameter(self._eet, self._youngs, 2)
        dx, dy = np.meshgrid(
            np.arange(self._grid.number_of_node_columns) * self._grid.dx,
            np.arange(self._grid.number_of_node_rows) * self._grid.dy)

        return kei(np.sqrt(dx ** 2 + dy ** 2) / alpha)

    def update(self, n_procs=1):
        elevation = self._nodal_values['lithosphere__elevation']
        load = self._nodal_values['lithosphere__overlying_pressure']
        deflection = self._nodal_values['lithosphere__elevation_increment']
        deposition = self._nodal_values['planet_surface_sediment__deposition_increment']

        new_load = ((load - self._last_load) +
                    (deposition * 2650. * 9.81).flat)

        self._last_load = load.copy()

        deflection.fill(0.)

        if self._method == 'airy':
            deflection[:] = new_load / (3300. * 9.81)
        else:
            self.subside_loads(new_load, deflection=deflection,
                               n_procs=n_procs)
        np.subtract(elevation, deflection, out=elevation)

    def subside_loads(self, loads, deflection=None, n_procs=1):
        if deflection is None:
            deflection = np.empty(self.shape, dtype=np.float)

        from .cfuncs import subside_grid_in_parallel

        w = deflection.reshape(self._grid.shape)
        load = loads.reshape(self._grid.shape)
        alpha = get_flexure_parameter(self._eet, self._youngs, 2)

        subside_grid_in_parallel(w, load, self._r, alpha, n_procs)

        return deflection


if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Changed __init__ method to accept file or keywords.<commit_after>#!/usr/bin/env python
"""Deform the lithosphere with 1D or 2D flexure.

Landlab component that implements a 1 and 2D lithospheric flexure
model.

Examples
--------

Create a grid on which we will run the flexure calculations.

>>> from landlab import RasterModelGrid
>>> from landlab.components.flexure import Flexure
>>> grid = RasterModelGrid((5, 4), spacing=(1.e4, 1.e4))

Check the fields that are used as input to the flexure component.

>>> Flexure.input_var_names
('lithosphere__elevation', 'lithosphere__overlying_pressure',
 'planet_surface_sediment__deposition_increment')

Check the units for the fields.

>>> Flexure.units['lithosphere__elevation']
'm'
>>> Flexure.units['lithosphere__overlying_pressure']
'Pa'

If you are not sure about one of the input or output variables, you can
get help for specific variables.

>>> Flexure.var_help('planet_surface_sediment__deposition_increment')

Add these fields to the grid.

>>> grid.add_zeros('lithosphere__elevation', at='node')
>>> grid.add_zeros('lithosphere__overlying_pressure', at='node')
>>> grid.add_zeros('planet_surface_sediment__deposition_increment', at='node')

>>> dh = grid.at_node['planet_surface_sediment__deposition_increment']
>>> dh.reshape(grid.shape)
>>> dh[2, 1:3] = 1e3

>>> flex = Flexure(grid)
>>> flex.update()

>>> flex.output_var_names
('lithosphere__elevation', 'lithosphere__elevation_increment')
>>> flex.grid.at_node['lithosphere__elevation']
"""

import numpy as np

from landlab import Component
from .funcs import get_flexure_parameter
from ...utils.decorators import use_file_name_or_kwds


class Flexure(Component):

    """Deform the lithosphere with 1D or 2D flexure.

    Landlab component that implements a 1 and 2D lithospheric flexure
    model.

    Examples
    --------
    >>> from landlab import RasterModelGrid
    >>> from landlab.components.flexure import Flexure
    >>> grid = RasterModelGrid((5, 4), spacing=(1.e4, 1.e4))

    >>> flex = Flexure(grid)
    >>> flex.name
    'Flexure'
    >>> sorted(flex.input_var_names) # doctest: +NORMALIZE_WHITESPACE
    ['lithosphere__elevation', 'lithosphere__overlying_pressure',
     'planet_surface_sediment__deposition_increment']
    >>> sorted(flex.output_var_names)
    ['lithosphere__elevation', 'lithosphere__elevation_increment']
    >>> for var in sorted(flex.units): flex.units[var]
    'm'
    'm'
    'Pa'
    'm'

    >>> flex.grid.number_of_node_rows
    5
    >>> flex.grid.number_of_node_columns
    4
    >>> flex.grid is grid
    True

    >>> np.all(grid.at_node['lithosphere__elevation'] == 0.)
    True

    >>> np.all(grid.at_node['lithosphere__elevation'] == 0.)
    True
    >>> np.all(grid.at_node['lithosphere__overlying_pressure'] == 0.)
    True
    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere__elevation_increment'] == 0.)
    True

    >>> load = grid.at_node['lithosphere__overlying_pressure']
    >>> load[4] = 1e9
    >>> dz = grid.at_node['lithosphere__elevation_increment']
    >>> np.all(dz == 0.)
    True

    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere__elevation_increment'] == 0.)
    False
    """

    _name = 'Flexure'

    _input_var_names = (
        'lithosphere__overlying_pressure',
        'lithosphere__elevation',
        'planet_surface_sediment__deposition_increment',
    )

    _output_var_names = (
        'lithosphere__elevation_increment',
        'lithosphere__elevation',
    )

    _var_units = {
        'lithosphere__overlying_pressure': 'Pa',
        'lithosphere__elevation': 'm',
        'lithosphere__elevation_increment': 'm',
        'planet_surface_sediment__deposition_increment': 'm',
    }

    _var_mapping = {
        'lithosphere__overlying_pressure': 'node',
        'lithosphere__elevation': 'node',
        'lithosphere__elevation_increment': 'node',
        'planet_surface_sediment__deposition_increment': 'node',
    }

    _var_doc = {
        'lithosphere__overlying_pressure':
            'The pressure at the base of the lithosphere',
        'lithosphere__elevation':
            'The elevation of the top of the lithosphere, i.e., the land '
            'surface',
        'lithosphere__elevation_increment':
            'The change in elevation of the top of the lithosphere (the land '
            'surface) in one timestep',
        'planet_surface_sediment__deposition_increment':
            'The amount of sediment deposited at the land surface in one '
            'timestep',
    }

    @use_file_name_or_kwds
    def __init__(self, grid, eet=65e3, youngs=7e10, method='airy'):
        """Initialize the flexure component.

        Parameters
        ----------
        grid : RasterModelGrid
            A grid.
        eet : float, optional
            Effective elastic thickness (m).
        youngs : float, optional
            Young's modulus.
        method : {'airy', 'flexure'}, optional
            Method to use to calculate deflections.
        """
        if method not in ('airy', 'flexure'):
            raise ValueError(
                '{method}: method not understood'.format(method=method))

        self._eet = eet
        self._youngs = youngs
        self._method = method

        self._grid = grid

        super(Flexure, self).__init__(grid)

        for name in self._input_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros('node', name, units=self._var_units[name])

        for name in self._output_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros('node', name, units=self._var_units[name])

        self._last_load = (
            self.grid.at_node['lithosphere__overlying_pressure'].copy())

        self._r = self._set_kei_func_grid()

    def _set_kei_func_grid(self):
        from scipy.special import kei

        alpha = get_flexure_parameter(self._eet, self._youngs, 2)
        dx, dy = np.meshgrid(
            np.arange(self._grid.number_of_node_columns) * self._grid.dx,
            np.arange(self._grid.number_of_node_rows) * self._grid.dy)

        return kei(np.sqrt(dx ** 2 + dy ** 2) / alpha)

    def update(self, n_procs=1):
        elevation = self.grid.at_node['lithosphere__elevation']
        load = self.grid.at_node['lithosphere__overlying_pressure']
        deflection = self.grid.at_node['lithosphere__elevation_increment']
        deposition = self.grid.at_node['planet_surface_sediment__deposition_increment']

        new_load = ((load - self._last_load) +
                    (deposition * 2650. * 9.81).flat)

        self._last_load = load.copy()

        deflection.fill(0.)

        if self._method == 'airy':
            deflection[:] = new_load / (3300. * 9.81)
        else:
            self.subside_loads(new_load, deflection=deflection,
                               n_procs=n_procs)
        np.subtract(elevation, deflection, out=elevation)

    def subside_loads(self, loads, deflection=None, n_procs=1):
        if deflection is None:
            deflection = np.empty(self.shape, dtype=np.float)

        from .cfuncs import subside_grid_in_parallel

        w = deflection.reshape(self._grid.shape)
        load = loads.reshape(self._grid.shape)
        alpha = get_flexure_parameter(self._eet, self._youngs, 2)

        subside_grid_in_parallel(w, load, self._r, alpha, n_procs)

        return deflection

{'fraction_non_alphanumeric': 0.11519865461425267, 'fraction_numerical': 0.007778011351692243, 'mean_word_length': 2.949086884338683, 'pattern_counts': {'":': 1, '<': 3, '<?xml version=': 0, '>': 190, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7042731', 'n_tokens_mistral': 4817, 'n_tokens_neox': 4434, 'n_words': 1000}