# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


from typing import Any
from epi_forecast_stat_mech.statistical_models import tree_util
import dataclasses
import jax
import jax.numpy as jnp
  Returns:
    The adjustment to the log likelihood due to BIC.
  """
  return - jnp.log(num_data_points) * num_parameters / 2


@dataclasses.dataclass
class BICModel:
  """A wrapper that applies Bayesian Information Criterion to a model."""

  base_model: Any
  bic_multiplier: float = 1.
  nonzero_threshold: float = .1
  nonzero_sharpness: float = 20

  def log_prior(self, parameters):
    return self.base_model.log_prior(parameters)

  def _count_nonzero(self, params):
    param_list, _ = tree_util.tree_flatten(params)
    return sum([
        soft_nonzero(p, self.nonzero_sharpness, self.nonzero_threshold).sum()
        for p in param_list])

  def log_likelihood(self, parameters, covariates, mechanistic_parameters):
    """Returns the log likelihood of `mechanistic_parameters`."""
    base_log_likelihood = self.base_model.log_likelihood(parameters,
                                                         covariates,
                                                         mechanistic_parameters)
    # `covariate_matrix` has shape `[num_epidemics, covariate_dim]`.

# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


# Path: epi_forecast_stat_mech/statistical_models/bic_model.py
from typing import Any
from epi_forecast_stat_mech.statistical_models import tree_util
import dataclasses
import jax
import jax.numpy as jnp

  Returns:
    The adjustment to the log likelihood due to BIC.
  """
  return - jnp.log(num_data_points) * num_parameters / 2


@dataclasses.dataclass
class BICModel:
  """A wrapper that applies Bayesian Information Criterion to a model."""

  base_model: Any
  bic_multiplier: float = 1.
  nonzero_threshold: float = .1
  nonzero_sharpness: float = 20

  def log_prior(self, parameters):
    return self.base_model.log_prior(parameters)

  def _count_nonzero(self, params):
    param_list, _ = tree_util.tree_flatten(params)
    return sum([
        soft_nonzero(p, self.nonzero_sharpness, self.nonzero_threshold).sum()
        for p in param_list])

  def log_likelihood(self, parameters, covariates, mechanistic_parameters):
    """Returns the log likelihood of `mechanistic_parameters`."""
    base_log_likelihood = self.base_model.log_likelihood(parameters,
                                                         covariates,
                                                         mechanistic_parameters)
    # `covariate_matrix` has shape `[num_epidemics, covariate_dim]`.

    covariate_matrix, _ = tree_util.pack(covariates)