Integration of DAE module

examples/ODEs/Part_6_NetworkODE.py

@@ -1,4 +1,4 @@
-# Numpy + plotting utilities + ordered dicts
+# %% Numpy + plotting utilities + ordered dicts
 import numpy as np
 import matplotlib.pyplot as plt
 from collections import OrderedDict
@@ -11,7 +11,7 @@
 
 # Neuromancer imports
 from neuromancer.psl.coupled_systems import *
-from neuromancer.dynamics import integrators, ode, physics, interpolation
+from neuromancer.dynamics import integrators, ode, physics
 from neuromancer.dataset import DictDataset
 from neuromancer.constraint import variable
 from neuromancer.problem import Problem
@@ -24,7 +24,7 @@
 np.random.seed(200)
 torch.manual_seed(0)
 
-
+# Define Network and datasets
 adj = np.array([[0,1],[0,2],[0,3],[1,0],[1,3],[1,4],[2,0],[2,3],[3,0],[3,1],[3,2],[3,4],[4,1],[4,3]]).T
 s = RC_Network(nx=5, adj=adj)
 nsim = 500
@@ -48,42 +48,46 @@
 train_dataset, dev_dataset, = [DictDataset(d, name=n) for d, n in zip([train_data, dev_data], ['train', 'dev'])]
 train_loader, dev_loader, test_loader = [DataLoader(d, batch_size=nsim//nstep, collate_fn=d.collate_fn, shuffle=True) for d in [train_dataset, dev_dataset, dev_dataset]]
 
-zones = [physics.RCNode(C=nn.Parameter(torch.tensor(5.0)),scaling=1.0e-5) for i in range(5)]  # heterogeneous population w/ identical physics
 
-heaters = [physics.SourceSink() for i in range(5)] # define heaters
+# Define the states
+states = {}
+states['T_1'] = 0
+states['T_2'] = 1
+states['T_3'] = 2
+states['T_4'] = 3
+states['T_5'] = 4
+states['T_6'] = 5
+states['T_7'] = 6
+states['T_8'] = 7
+states['T_9'] = 8
+states['T_10'] = 9
+states['T_11'] = 10
+
+# Model construction
+keys = list(states.keys())
+zones = [physics.RCNode(in_keys=[keys[i]], state_keys=[keys[i]], 
+                        C=nn.Parameter(torch.tensor(5.0)),scaling=1.0e-5) for i in range(5)]
 
-outside = [physics.SourceSink()]
+heaters = [physics.SourceSink(state_keys=[keys[i+len(zones)]], in_keys=[keys[i+len(zones)]]) for i in range(5)] # define heaters
+
+outside = [physics.SourceSink(state_keys=[keys[-1]], in_keys=[keys[-1]])]
 
 # join lists:
 agents = zones + heaters + outside
 
-map = physics.map_from_agents(agents)
-# Let's take a look at this 'map':
-print(map)
-
 # Helper function for constructing couplings based on desired edge physics and an edge list:
-def generate_parameterized_edges(physics,edge_list):
+def generate_deltaTemp_edges(physics,edge_list,agents):
     """
     Quick helper function to construct edge physics/objects from adj. list:
     """
-
     couplings = []
-    if isinstance(physics,nn.Module): # is "physics" an instance or a class?
-        # If we're in here, we expect one instance of "physics" for all edges in edge_list (homogeneous edges)
-        physics.pins = edge_list
-        couplings.append(physics)
-        print(f'Broadcasting {physics} to all elements in edge list.')
-    else:
-        # If we're in here, we expect different "physics" for each edge in edge_list (heterogeneous edges)
-        for edge in edge_list:
-            agent = physics(R=nn.Parameter(torch.tensor(50.0)),pins=[edge])
-            couplings.append(agent)
-
-        print(f'Assuming new {physics} for each element in edge list.')
+    for edge in edge_list:
+        agent = physics(in_keys=[*agents[edge[1]].in_keys,*agents[edge[0]].in_keys],R=nn.Parameter(torch.tensor(50.0)),pins=[edge])
+        couplings.append(agent)
 
     return couplings
 
-couplings = generate_parameterized_edges(physics.DeltaTemp,list(adj.T))    # Heterogeneous edges of same physics
+couplings = generate_deltaTemp_edges(physics.DeltaTemp,list(adj.T),agents)    # Heterogeneous edges of same physics
 
 # What do we have so far?
 print(len(couplings))
@@ -93,28 +97,25 @@ def generate_parameterized_edges(physics,edge_list):
 print(couplings[0].pins)
 
 # Couple w/ outside temp:
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[0,5]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[1,5]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[2,5]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[3,5]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[4,5]]))
+outside_list = [[0,5],[1,5],[2,5],[3,5],[4,5]]
+out_couplings = generate_deltaTemp_edges(physics.DeltaTemp,outside_list,agents)
 
 # Couple w/ individual sources:
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[0,6]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[1,7]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[2,8]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[3,9]]))
-couplings.append(physics.DeltaTemp(R=nn.Parameter(torch.tensor(50.0)),pins=[[4,10]]))
+source_list = [[0,6],[1,7],[2,8],[3,9],[4,10]]
+source_couplings = generate_deltaTemp_edges(physics.DeltaTemp,source_list,agents)
+
+couplings += out_couplings + source_couplings
 
+# Model ODE RHS instantiation
 model_ode = ode.GeneralNetworkedODE(
-    map = map,
+    states=states,
     agents = agents,
     couplings = couplings,
     insize = s.nx+s.nu,
-    outsize = s.nx,
-    inductive_bias="compositional")
+    outsize = s.nx)
 
-fx_int = integrators.RK2(model_ode, h=1.0)
+# Integrator instantiation
+fx_int = integrators.RK4(model_ode, h=1.0)
 
 dynamics_model = System([Node(fx_int,['xn','U'],['xn'])])
 
@@ -167,4 +168,5 @@ def generate_parameterized_edges(physics,edge_list):
 
 plt.figure()
 plt.plot(sol.detach().numpy(),label='model', color = 'black')
-plt.plot(s_test['X'][:,:5],label = 'data', color = 'red')
+plt.plot(s_test['X'][:,:5],label = 'data', color = 'red')
+# %%