code: def partially_observed_lorenz_system(
      observed_values,
      innovation_scale,
      observation_scale,
      observation_mask,
      step_size,
      observation_index,
  ):
      """Lorenz System model.
  
      Args:
        observed_values: Array of observed values.
        innovation_scale: Python `float`.
        observation_scale: Python `float`.
        observation_mask: `bool` array used to occlude observations.
        step_size: Python `float`.
        observation_index: `int` index used to pick which latent time series
          is observed.
  
      Returns:
        model: `StanModel`.
      """
  
      code = (
          """
    data {
      int<lower=0> num_timesteps;
      int<lower=0> num_observations;
      int<lower=1> observation_state_index;
      int<lower = 1, upper = num_timesteps> observation_time_indices[
        num_observations];
      vector[num_observations] observations;
      real<lower=0> innovation_scale;
      real<lower=0> observation_scale;
      real<lower=0> step_size;
    }
    parameters {
      matrix[num_timesteps, 3] latents;
    }
    model {"""
          + LORENZ_SYSTEM_OBSERVATIONS_MODEL
          + "}"
      )
  
      stan_data = {
          "num_timesteps": len(observed_values),
          "num_observations": len(observed_values[observation_mask]),
          "observation_time_indices": np.arange(
              1, len(observed_values) + 1
          )[observation_mask],
          "observations": observed_values[observation_mask],
          "observation_state_index": observation_index + 1,
          "innovation_scale": innovation_scale,
          "observation_scale": observation_scale,
          "step_size": step_size,
      }
  
      model = util.cached_stan_model(code)
  
      def _ext_identity(samples):
          """Extracts the values of all latent variables."""
          latents = util.get_columns(samples, r"^latents\[\d+,\d+\]$")
          # Last two dimensions are swapped in Stan output.
          return latents.reshape((-1, 3, 30)).swapaxes(1, 2)
  
      extract_fns = {"identity": _ext_identity}
  
      return stan_model.StanModel(
          extract_fns=extract_fns,
          sample_fn=util.make_sample_fn(model, data=stan_data),
      )
  
imports: import numpy
  from inference_gym.tools.stan import stan_model
  from inference_gym.tools.stan import util
label_desc: Use image node_id 1 for calling a global function with example usage: partially_observed_lorenz_system(observed_values, innovation_scale, observation_scale, observation_mask, step_size, observation_index) and returns: stan_model, latents