Project: s41467-024-51477-5 | Function: make_fourier_features (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the spatial Fourier features generation from Equation 9, creating sine and cosine basis functions with exponentially increasing frequencies to capture high-frequency spatial patterns

**KEY TERMS:** {
  "degrees": "harmonic indices h∈H_i^s",
  "2**degrees": "exponential frequency scaling 2^h",
  "y": "2π2^h·x computation",
  "x": "spatial coordinate s_i",
  "jnp.cos(y), jnp.sin(y)": "cos(2π2^h s_i), sin(2π2^h s_i)"
}

**CONTEXT CODE:** from jax import numpy as jnp
import numpy as np
from tensorflow_probability.substrates import jax as tfp  # pylint: disable=g-importing-member

tfd = tfp.distributions


class LikelihoodDist(enum.Enum):
  NORMAL = 'NORMAL'
  NB = 'NB'
  ZINB = 'ZINB'


def make_seasonal_frequencies(
    seasonality_periods: np.ndarray, num_harmonics: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
  """Return unique Fourier frequencies for given periods and harmonics."""
  seasonality_periods = np.array(seasonality_periods, dtype=np.float32)
  if np.any((num_harmonics > seasonality_periods / 2)):
    raise ValueError('Harmonic cannot exceed half seasonal period.')
  if seasonality_periods.shape != num_harmonics.shape:
    raise ValueError('Number of seasonal periods and harmonics must be equal.')
  if len(num_harmonics.shape) != 1:
    raise ValueError(
        'Arguments `num_harmonics` and `seasonality_periods` must be rank 1.'
    )
  if seasonality_periods.shape[0] == 0:
    return (np.zeros(0), np.zeros(0))
  harmonics = [np.arange(1, h + 1, dtype=np.float32) for h in num_harmonics]
  frequencies = np.concatenate(
      [h / p for (h, p) in zip(harmonics, seasonality_periods)]
  )
  _, idx = np.unique(frequencies, return_index=True)
  idx_sort = np.sort(idx)
  unique_frequencies = frequencies[idx_sort]
  unique_harmonics = np.concatenate(harmonics)[idx_sort]
  return (unique_frequencies, unique_harmonics)


def make_seasonal_features(
    x: jax.typing.ArrayLike,
    seasonality_periods: np.ndarray,
    num_harmonics: np.ndarray,
    rescale: bool = False,
) -> jnp.ndarray:
  """Returns a set of cos and sin features for each seasonality period."""
  # MASKED: Implements the temporal seasonal features generation described in Equation 8 of the paper, creating cosine and sine features for each seasonality period using harmonic frequencies
  raise NotImplementedError("Function has been masked for testing purposes")


def make_fourier_features(
    x: jax.typing.ArrayLike, max_degree: int, rescale: bool = False
) -> jnp.ndarray:
  """Returns a set of sine and cosine basis functions."""
  # MASKED: Implements the spatial Fourier features generation from Equation 9, creating sine and cosine basis functions with exponentially increasing frequencies to capture high-frequency spatial patterns
  raise NotImplementedError("Function has been masked for testing purposes")


prior_base_d = tfd.Logistic


def prior_model_fn(mlp_template):
  yield prior_base_d(0.0, 1.0)  # log_noise_scale
  yield prior_base_d(-1.5, 1.0)  # shape
  yield prior_base_d(0.0, 1.0)  # inflated_loc_probs

  leaves = jax.tree_util.tree_leaves(mlp_template)
  flat_mlp_params = []
  for p in leaves:
    weights = yield prior_base_d(0.0, jnp.ones_like(p))
    flat_mlp_params.append(weights)


def make_likelihood_model(
    params: jax.typing.ArrayLike,
    x: jax.Array,
    mlp: nn.Module,
    mlp_template: flax.core.frozen_dict.FrozenDict,
    distribution: str,
) -> tfd.Distribution:
  """Creates likelihood for the model.

  For a Gaussian likelihood, the likelihood is centered at the predictions and
  the scale is given by the noise scale.

  For a NegativeBinomial, we parametrize in terms of mean and shape
  following [DeepAR: Probabilistic Forecasting with
  Autoregressive Recurrent Networks](https://arxiv.org/pdf/1704.04110.pdf)
  (Salinas et al. 2019; Section 3.1)
  but use a shared shape parameter, which is the inverse of the
  total count. The variance equals `mean + mean^2 * shape`.

  See also <https://en.wikipedia.org/wiki/Negative_binomial_distribution>, in
  particular the section on Alternative parameterization, which defines:

  ```
  probs = mean / variance
  total_count = mean^2 / (variance - mean)
  ```

  so that `variance = mean + mean^2 / total_count`
  and thus `total_count = 1 / shape`.

  From this we get the parameters passed to `tfd.NegativeBinomial` below:

  ```
  total_count = 1 / shape
  probs = mean / (mean + mean^2 * shape)
        = 1 / (1 + mean * shape)
  ```

  For zero inflation we follow ideas from https://arxiv.org/pdf/2010.09647.pdf

  Args:
    params: Pytree of particles.
    x: Time indices for the training set.
    mlp: BayesianNeuralField1D model instance.
    mlp_template: MLP variables.
    distribution: Either 'NORMAL', 'NB' (Negative Binomial) or 'ZINB' (Zero
      Inflated Negative Binomial) are supported.

  Returns:
    A distribution with the likelihood for each sample in the batch.
  """
  if LikelihoodDist(distribution) == LikelihoodDist.NORMAL:
    log_noise_scale = params[0]
    treedef = jax.tree_util.tree_structure(mlp_template)
    mlp_params = jax.tree_util.tree_unflatten(treedef, params[3:])
    predictions = mlp.apply(mlp_params, x)
    return tfd.Independent(
        tfd.Normal(predictions, 0.01 + jnp.exp(log_noise_scale)), 1
    )

  elif LikelihoodDist(distribution) == LikelihoodDist.NB:
    treedef = jax.tree_util.tree_structure(mlp_template)
    mlp_params = jax.tree_util.tree_unflatten(treedef, params[3:])
    predictions = mlp.apply(mlp_params, x)
    mean = jax.nn.softplus(predictions)
    shape = jax.nn.softplus(params[1])

    negative_binomial = tfd.NegativeBinomial(
        total_count=1 / shape, logits=-jnp.log(shape) - jnp.log(mean)
    )
    return tfd.Independent(negative_binomial, 1)

  elif LikelihoodDist(distribution) == LikelihoodDist.ZINB:
    treedef = jax.tree_util.tree_structure(mlp_template)
    mlp_params = jax.tree_util.tree_unflatten(treedef, params[3:])
    predictions = mlp.apply(mlp_params, x)
    mean = jax.nn.softplus(predictions)


Complete the function implementation replacing "raise NotImplementedError".