allow users to pass 1dim conditions

Author: radiradev

test_condtional.png renamed to assets/test_condtional.png

@@ -94,6 +94,8 @@ def train_single(self, xt, vt, conditions=None):
         )
 
         if conditions is not None:
+            if conditions.ndim == 1:
+                conditions = np.expand_dims(conditions, axis=1)
             xt = np.concatenate([xt, conditions], axis=1)
 
         # learn to predict the velocity field given a noised input
@@ -158,6 +160,8 @@ def model_t(self, t, xt, conditions=None):
         flow_step = int(round(t * (self.n_flow_steps - 1)))
 
         if conditions is not None:
+            if conditions.ndim == 1:
+                conditions = np.expand_dims(conditions, axis=1)
             xt = np.concatenate([xt, conditions], axis=1)
 
         return self.models[flow_step].predict(xt)

flowmatching_bdt/flow_bdt.py

@@ -25,7 +25,7 @@ def test_condtional():
 
     # get new samples
     num_samples = 1000
-    conditions = np.ones((num_samples, 1))
+    conditions = np.ones(num_samples)
     samples = model.predict(num_samples=num_samples, conditions=conditions)
 
 

flowmatching_bdt/test/test_model.py

@@ -3,7 +3,7 @@
 setup(
   name = 'flowmatching-bdt',
   packages = find_packages(exclude=['assets']),
-  version = '0.1.0',
+  version = '0.2.0',
   license='MIT',
   description = 'Flow Matching with BDTs',
   long_description_content_type = 'text/markdown',