Skip to content

Commit

Permalink
Draft of coupling class
Browse files Browse the repository at this point in the history
  • Loading branch information
mdbartos committed Mar 27, 2024
1 parent b4c598c commit c51df7e
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 2 deletions.
32 changes: 32 additions & 0 deletions pipedream_solver/coupling.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
class SpeciesCoupling():
def __init__(self, species={}, functions=[], parameters={}):
self.species = {}
self.functions = []
self.parameters = {}
for key, value in species.items():
if hasattr(self, key):
raise NameError
setattr(self, key, value)
self.species[key] = value
for function in functions:
if hasattr(self, function.__name__):
raise NameError
setattr(self, function.__name__, function)
self.functions.append(function)
for key, value in parameters.items():
if hasattr(self, key):
raise NameError
setattr(self, key, value)
self.parameters[key] = value

def step(self, **kwargs):
'''
Steps through all functions in `self.functions` in order.
'''
for function in self.functions:
arg_keys = function.__code__.co_varnames
# NOTE: Has potential to fail silently
arg_values = (getattr(self, key, None) for key in arg_keys)
function_inputs = dict(zip(arg_keys, arg_values))
function_inputs.update(kwargs)
result = function(**function_inputs)
24 changes: 22 additions & 2 deletions pipedream_solver/nquality.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ def __init__(self, hydraulics, superjunction_params, superlink_params,
self.N_process_sigma = superlink_params['KF_process_sigma'].values[0].astype(np.float64)
self.N_measure_sigma = superlink_params['KF_measure_sigma'].values[0].astype(np.float64)
self.step(dt=1e-6)

# TODO: It might be safer to have these as @properties
def import_hydraulic_states(self, _dt):
self._H_j_next = self.hydraulics.H_j
Expand Down Expand Up @@ -277,7 +277,27 @@ def import_hydraulic_states(self, _dt):
self.hydraulics._geom_codes, self.hydraulics._g1_ik, self._H_j_next, self._z_inv_j, self._H_j_prev, self._Q_dk_next, self._B_dk, self._dx_dk,
self._Q_uk_up_next, self._Q_uk_dn_next, self._Q_dk_up_next, self._Q_dk_dn_next,
self._A_uk_next, self._A_dk_next, self._A_uk_prev, self._A_dk_prev)


@property
def c(self):
# NOTE: more performant to write in-place
result = np.concatenate([self._c_j, self._c_Ik, self._c_ik,
self._c_uk, self._c_dk])
return result

@c.setter
def c(self, value):
n1 = self._c_j.size
n2 = n1 + self._c_Ik.size
n3 = n2 + self._c_ik.size
n4 = n3 + self._c_uk.size
n5 = n4 + self._c_dk.size
self._c_j[:] = np.asarray(value[:n1])
self._c_Ik[:] = np.asarray(value[n1:n2])
self._c_ik[:] = np.asarray(value[n2:n3])
self._c_uk[:] = np.asarray(value[n3:n4])
self._c_dk[:] = np.asarray(value[n4:n5])

@property
def c_j(self):
return self._c_j
Expand Down

0 comments on commit c51df7e

Please sign in to comment.