<gh_stars>0
"""
The code for the mean sea level model
"""

from collections import OrderedDict
import multiprocessing
import os
from typing import Tuple

from numba import jitclass, float64
import pandas as pd

from .util import compile_model, hash_string, DATA_DIR, CACHE_DIR, gev_cdf, random_gev
from .param import FloodPriorParameters


@jitclass(
    OrderedDict(
        loc_base=float64,
        loc_trend=float64,
        coeff_var=float64,
        shape=float64,
        zero_time=float64,
        scale_min=float64,
    )
)
class FloodModel:
    """
    This is the state of world, or simulation-specific set of parameters
    governing 
    """

    def __init__(
        self,
        loc_base: float,
        loc_trend: float,
        coeff_var: float,
        shape: float,
        zero_time: float,
        scale_min: float,
    ) -> None:
        """
        Build the model
        """
        self.loc_base = loc_base
        self.loc_trend = loc_trend
        self.coeff_var = coeff_var
        self.shape = shape
        self.zero_time = zero_time
        self.scale_min = scale_min

    def get_gev_param(self, year: float) -> Tuple[float, float, float]:
        """
        Get the location, scale, and shape parameter for a particular year
        """
        year_adj = year - self.zero_time
        loc = self.loc_base + self.loc_trend * year_adj
        scale = max(loc * self.coeff_var, self.scale_min)
        shape = self.shape
        return loc, scale, shape

    def simulate_flood(self, year: float) -> float:
        """
        Simulate a storm surge for a particular year
        """
        loc, scale, shape = self.get_gev_param(year=year)
        return random_gev(loc=loc, scale=scale, shape=shape)

    def calc_exceedance_prob(self, year: float, height: float) -> float:
        """
        Calculate the probability that the flood in `year` is greater than `height`
        """
        loc, scale, shape = self.get_gev_param(year=year)
        return 1.0 - gev_cdf(x=height, loc=loc, scale=scale, shape=shape)


class FloodSimulationLibrary:
    """
    Stores a large number of simulations of the flood parameters
    """

    def __init__(self, param: FloodPriorParameters, **kwargs) -> None:
        """
        Build a library of simulation parameters
        """

        self.param = param
        simulation_fname = self.get_filename()

        try:
            self.simulations = pd.read_feather(simulation_fname)

        except IOError:
            self.simulations = self.create_simulation_library(**kwargs)
            self.simulations.to_feather(simulation_fname)

    def get_filename(self) -> str:
        """
        Get a unique filename for the library where the resulting library
        can be stored
        """
        simulation_hash = hash_string(OrderedDict(self.param._asdict()).__str__())
        return os.path.join(CACHE_DIR, f"{simulation_hash}.feather")

    def create_simulation_library(
        self,
        n_sim: int = 100_000,
        n_chains: int = multiprocessing.cpu_count(),
        n_warmup: int = 2500,
    ) -> pd.DataFrame:
        """
        Run the quadratic model to create many simulations of
        the SeaLevelModel parameters
        """
        sea_level = pd.read_csv(
            os.path.join(DATA_DIR, "sea-level-annual.csv"), index_col="year"
        )
        stan_fname = os.path.join(DATA_DIR, "gev-trend.stan")
        model = compile_model(filename=stan_fname, model_name="stormsurge")
        stan_data = dict(
            N=sea_level.shape[0],
            y=sea_level["storm_surge"].values,
            time=sea_level.index.values,
            # PRIORS
            time_offset=self.param.zero_time,
            loc_base_expected=self.param.loc_base_expected,
            loc_base_std=self.param.loc_base_std,
            loc_trend_expected=self.param.loc_trend_expected,
            loc_trend_std=self.param.loc_trend_std,
            coeff_var_expected=self.param.coeff_var_expected,
            coeff_var_std=self.param.coeff_var_std,
            sigma_min=self.param.sigma_min,
        )

        # calculate the number of iterations needed
        n_iter = n_sim / n_chains + n_warmup
        fitted = model.sampling(
            data=stan_data, chains=int(n_chains), iter=int(n_iter), warmup=int(n_warmup)
        )
        return pd.DataFrame(
            fitted.extract(["loc_base", "loc_trend", "coeff_var", "shape"])
        )

    def build_model(self) -> FloodModel:
        """
        Get a model with randomly sampled state of the world
        """
        row = self.simulations.sample(1, axis=0)
        return FloodModel(
            zero_time=self.param.zero_time,
            loc_base=row["loc_base"].values[0],
            loc_trend=row["loc_trend"].values[0],
            coeff_var=row["coeff_var"].values[0],
            shape=row["shape"].values[0],
            scale_min=self.param.sigma_min,
        )

"""
The code for the mean sea level model
"""

from collections import OrderedDict
import multiprocessing
import os
from typing import Tuple

from numba import jitclass, float64
import pandas as pd

from .util import compile_model, hash_string, DATA_DIR, CACHE_DIR, gev_cdf, random_gev
from .param import FloodPriorParameters


@jitclass(
    OrderedDict(
        loc_base=float64,
        loc_trend=float64,
        coeff_var=float64,
        shape=float64,
        zero_time=float64,
        scale_min=float64,
    )
)
class FloodModel:
    """
    This is the state of world, or simulation-specific set of parameters
    governing 
    """

    def __init__(
        self,
        loc_base: float,
        loc_trend: float,
        coeff_var: float,
        shape: float,
        zero_time: float,
        scale_min: float,
    ) -> None:
        """
        Build the model
        """
        self.loc_base = loc_base
        self.loc_trend = loc_trend
        self.coeff_var = coeff_var
        self.shape = shape
        self.zero_time = zero_time
        self.scale_min = scale_min

    def get_gev_param(self, year: float) -> Tuple[float, float, float]:
        """
        Get the location, scale, and shape parameter for a particular year
        """
        year_adj = year - self.zero_time
        loc = self.loc_base + self.loc_trend * year_adj
        scale = max(loc * self.coeff_var, self.scale_min)
        shape = self.shape
        return loc, scale, shape

    def simulate_flood(self, year: float) -> float:
        """
        Simulate a storm surge for a particular year
        """
        loc, scale, shape = self.get_gev_param(year=year)
        return random_gev(loc=loc, scale=scale, shape=shape)

    def calc_exceedance_prob(self, year: float, height: float) -> float:
        """
        Calculate the probability that the flood in `year` is greater than `height`
        """
        loc, scale, shape = self.get_gev_param(year=year)
        return 1.0 - gev_cdf(x=height, loc=loc, scale=scale, shape=shape)


class FloodSimulationLibrary:
    """
    Stores a large number of simulations of the flood parameters
    """

    def __init__(self, param: FloodPriorParameters, **kwargs) -> None:
        """
        Build a library of simulation parameters
        """

        self.param = param
        simulation_fname = self.get_filename()

        try:
            self.simulations = pd.read_feather(simulation_fname)

        except IOError:
            self.simulations = self.create_simulation_library(**kwargs)
            self.simulations.to_feather(simulation_fname)

    def get_filename(self) -> str:
        """
        Get a unique filename for the library where the resulting library
        can be stored
        """
        simulation_hash = hash_string(OrderedDict(self.param._asdict()).__str__())
        return os.path.join(CACHE_DIR, f"{simulation_hash}.feather")

    def create_simulation_library(
        self,
        n_sim: int = 100_000,
        n_chains: int = multiprocessing.cpu_count(),
        n_warmup: int = 2500,
    ) -> pd.DataFrame:
        """
        Run the quadratic model to create many simulations of
        the SeaLevelModel parameters
        """
        sea_level = pd.read_csv(
            os.path.join(DATA_DIR, "sea-level-annual.csv"), index_col="year"
        )
        stan_fname = os.path.join(DATA_DIR, "gev-trend.stan")
        model = compile_model(filename=stan_fname, model_name="stormsurge")
        stan_data = dict(
            N=sea_level.shape[0],
            y=sea_level["storm_surge"].values,
            time=sea_level.index.values,
            # PRIORS
            time_offset=self.param.zero_time,
            loc_base_expected=self.param.loc_base_expected,
            loc_base_std=self.param.loc_base_std,
            loc_trend_expected=self.param.loc_trend_expected,
            loc_trend_std=self.param.loc_trend_std,
            coeff_var_expected=self.param.coeff_var_expected,
            coeff_var_std=self.param.coeff_var_std,
            sigma_min=self.param.sigma_min,
        )

        # calculate the number of iterations needed
        n_iter = n_sim / n_chains + n_warmup
        fitted = model.sampling(
            data=stan_data, chains=int(n_chains), iter=int(n_iter), warmup=int(n_warmup)
        )
        return pd.DataFrame(
            fitted.extract(["loc_base", "loc_trend", "coeff_var", "shape"])
        )

    def build_model(self) -> FloodModel:
        """
        Get a model with randomly sampled state of the world
        """
        row = self.simulations.sample(1, axis=0)
        return FloodModel(
            zero_time=self.param.zero_time,
            loc_base=row["loc_base"].values[0],
            loc_trend=row["loc_trend"].values[0],
            coeff_var=row["coeff_var"].values[0],
            shape=row["shape"].values[0],
            scale_min=self.param.sigma_min,
        )
The code for the mean sea level model This is the state of world, or simulation-specific set of parameters     governing Build the model Get the location, scale, and shape parameter for a particular year Simulate a storm surge for a particular year Calculate the probability that the flood in `year` is greater than `height` Stores a large number of simulations of the flood parameters Build a library of simulation parameters Get a unique filename for the library where the resulting library         can be stored Run the quadratic model to create many simulations of         the SeaLevelModel parameters # PRIORS # calculate the number of iterations needed Get a model with randomly sampled state of the world