Flows without upstream stocks cause a KeyError in stockflow

[sabinala]

mira/notebooks/hackathon_2024.02/scenario5/scenario_5_stockflow.json

Line 15 in 17b194e

"upstream_stock": null,

MODEL_PATH = "https://raw.githubusercontent.com/gyorilab/mira/main/notebooks/hackathon_2024.02/scenario5/"
model = os.path.join(MODEL_PATH, "scenario_5_stockflow.json")
start_time = 1e-5
end_time = 300.0
logging_step_size = 10.0
num_samples = 1
result = pyciemss.sample(model, end_time, logging_step_size, num_samples, start_time=start_time)

results in this error

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[16], line 6
      3 logging_step_size = 10.0
      4 num_samples = 1
----> 6 result = pyciemss.sample(model, end_time, logging_step_size, num_samples, start_time=start_time)
      7 # display(result['data'].head())
      8 
      9 # Plot results
     10 schema = plots.trajectories(result["data"], keep=".*_state")

File ~/Projects/pyciemss/pyciemss/integration_utils/custom_decorators.py:29, in pyciemss_logging_wrapper.<locals>.wrapped(*args, **kwargs)
     17 log_message = """
     18     ###############################
     19 
   (...)
     26     ################################
     27 """
     28 logging.exception(log_message, function.__name__, function.__doc__)
---> 29 raise e

File ~/Projects/pyciemss/pyciemss/integration_utils/custom_decorators.py:10, in pyciemss_logging_wrapper.<locals>.wrapped(*args, **kwargs)
      8 try:
      9     start_time = time.perf_counter()
---> 10     result = function(*args, **kwargs)
     11     end_time = time.perf_counter()
     12     logging.info(
     13         "Elapsed time for %s: %f", function.__name__, end_time - start_time
     14     )

File ~/Projects/pyciemss/pyciemss/interfaces.py:239, in sample(model_path_or_json, end_time, logging_step_size, num_samples, noise_model, noise_model_kwargs, solver_method, solver_options, start_time, inferred_parameters, static_state_interventions, static_parameter_interventions, dynamic_state_interventions, dynamic_parameter_interventions)
    171 """
    172 Load a model from a file, compile it into a probabilistic program, and sample from it.
    173 
   (...)
    235                 and (num_samples,) for parameters.
    236 """
    238 with torch.no_grad():
--> 239     model = CompiledDynamics.load(model_path_or_json)
    241     logging_times = torch.arange(
    242         start_time + logging_step_size, end_time, logging_step_size
    243     )
    245     # Check that num_samples is a positive integer

File ~/anaconda3/lib/python3.10/functools.py:926, in singledispatchmethod.__get__.<locals>._method(*args, **kwargs)
    924 def _method(*args, **kwargs):
    925     method = self.dispatcher.dispatch(args[0].__class__)
--> 926     return method.__get__(obj, cls)(*args, **kwargs)

File ~/Projects/pyciemss/pyciemss/compiled_dynamics.py:108, in CompiledDynamics._load_from_url_or_path(cls, path)
    106     model = mira.sources.amr.model_from_url(path)
    107 else:
--> 108     model = mira.sources.amr.model_from_json_file(path)
    109 return cls.load(model)

File ~/anaconda3/lib/python3.10/site-packages/mira/sources/amr/__init__.py:45, in model_from_json_file(fname)
     43 with open(fname) as fh:
     44     model_json = json.load(fh)
---> 45 return model_from_json(model_json)

File ~/anaconda3/lib/python3.10/site-packages/mira/sources/amr/__init__.py:71, in model_from_json(model_json)
     69     return regnet.template_model_from_amr_json(model_json)
     70 elif 'stockflow' in header['schema']:
---> 71     return stockflow.template_model_from_amr_json(model_json)
     72 else:
     73     raise ValueError(f'Unknown schema: {header["schema"]}')

File ~/anaconda3/lib/python3.10/site-packages/mira/sources/amr/stockflow.py:112, in template_model_from_amr_json(model_json)
    105 # A stock is considered a controller if it has a link to the given
    106 # flow but is not an input to the flow, and is not an auxiliary
    107 controllers = [link['source'] for link in links if (
    108     link['target'] == flow['id']
    109     and link['source'] != input
    110     and link['source'] not in aux_expressions)]
--> 112 input_concepts = [concepts[input].copy(deep=True)]
    113 output_concepts = [concepts[output].copy(deep=True)]
    114 controller_concepts = [concepts[i].copy(deep=True) for i in controllers]

KeyError: None