Scale output, pinns method: DVM DBM (beta mode) (#13)

* first implementation * fixes * fixes * fixes * fixes * fixes * fixes * fixes * fixes * Gsolv direct * implemented DBM * fix * fix * fix * fix * fixes * fixes * Update Equations.py * m * DBM implemented * fixes * fixes * fixes * fixes * tests * fix * fix * fixes * fixes * Update test_xppbe.py * Update Equations.py * fix * fix * fix * bugfix * changes phi interface * fix * finally fixed * dG_dn * fix * fix * ready for testing * postcode * changes * m * order * order * order * m * m * M * m * m * m * m * m * m * m --------- Co-authored-by: martinachondo <[email protected]>
MartinAchondo · Jul 7, 2024 · 9815c0a · 9815c0a
1 parent af293e4
commit 9815c0a
Show file tree

Hide file tree

Showing 14 changed files with 726 additions and 271 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -20,7 +20,6 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install .
-        pip install .[pbe-standard]
         pip install pytest
 
     - name: Run tests

diff --git a/pyproject.toml b/pyproject.toml
@@ -18,9 +18,13 @@ dependencies = [
     "numpy==1.23.5",
     "tensorflow==2.10.0",
     "scipy==1.13.1",
+    "numba==0.59.1",
     "trimesh==3.23.5",
     "pygamer==2.0.7",
+    "meshio==5.3.4",
+    "gmsh==4.9.5",
     "rtree==1.0.1",
+    "bempp-cl==0.3.0",
     "matplotlib==3.5.1",
     "plotly==5.7.0",
     "kaleido==0.2.1",
@@ -33,12 +37,6 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-pbe-standard = [
-    "numba==0.59.1",
-    "meshio==5.3.4",
-    "gmsh==4.9.5",
-    "bempp-cl==0.3.0"
-]
 pbj = [
     "pbj@git+https://github.com/bem4solvation/pbj"
 ]

diff --git a/tests/simulations_yaml/test_born_ion.yaml b/tests/simulations_yaml/test_born_ion.yaml
@@ -55,6 +55,7 @@ hyperparameters_out:
     architecture_Net: FCNN
     fourier_features: true
     weight_factorization: true
+    scale_output: true
 
 
 # Solve parameters

diff --git a/tests/test_xppbe.py b/tests/test_xppbe.py
@@ -6,7 +6,7 @@
 from xppbe import Simulation
 from xppbe import Allrun,Allclean
 
-def run_checkers(sim,sim_name,temp_dir):
+def run_checkers(sim,sim_name,temp_dir,num_iters=1):
 
     results_path = os.path.join(temp_dir,'results')
     assert os.path.isdir(results_path)
@@ -17,12 +17,12 @@ def run_checkers(sim,sim_name,temp_dir):
         assert name in os.listdir(results_path)
 
     mesh_path = os.path.join(results_path,'mesh')
-    assert len(os.listdir(mesh_path)) == 10 or len(os.listdir(mesh_path)) == 11 
+    assert len(os.listdir(mesh_path)) > 0 
 
     iterations_path = os.path.join(results_path,'iterations')
-    assert len(os.listdir(iterations_path)) == 1
+    assert len(os.listdir(iterations_path)) == num_iters
 
-    last_iteration = os.path.join(iterations_path,f'iter_{sim.N_iters}')
+    last_iteration = os.path.join(iterations_path,f'iter_{int(sim.N_iters+sim.N_steps_2*sim.options_optimizer2["maxiter"])}')
     listdir_last_iteration = os.listdir(last_iteration)
 
     for file_name in ['optimizer.npy', 'weights.index', 'w_hist.csv', 'checkpoint', 'loss.csv', 'hyperparameters.json', 'G_solv.csv', 'loss_validation.csv', 'weights.data-00000-of-00001']:
@@ -33,7 +33,7 @@ def run_checkers(sim,sim_name,temp_dir):
         with open(file, 'r') as csvfile:
             reader = csv.reader(csvfile)
             lines = list(reader)
-            assert len(lines) == sim.N_iters + 1
+            assert len(lines) == int(sim.N_iters+sim.N_steps_2*sim.options_optimizer2["maxiter"]) + 1
             last_line = lines[-1]
             assert not '' in last_line
 
@@ -53,23 +53,35 @@ def test_scripts():
         run_checkers(sim,sim_name,temp_dir)
         Allclean(results_path=temp_dir)
         assert len(os.listdir(os.path.join(temp_dir,'results'))) == 0
-
+
+
 
 @pytest.mark.parametrize(
- ('molecule'),
+ ('pinns_method','model','scheme'),
  (
-     ('born_ion'),
-     ('sphere_+1-1'),
-     ('methanol'),
-     ('arg')
+     ('DCM','linear','direct'),
+     ('DCM','linear', 'regularized_scheme_1'),
+     ('DCM','nonlinear','regularized_scheme_2'),
+     ('DVM','linear','direct'),
+     ('DVM','linear','regularized_scheme_1'),
+     ('DVM','nonlinear','regularized_scheme_1'),
+     ('DBM','linear','direct')
  )       
 )
-def test_xppbe_solver(molecule):
+def test_pinns_method_and_schemes(pinns_method,model,scheme):
     with tempfile.TemporaryDirectory() as temp_dir:
-        sim_name = f'test_{molecule}'
+        sim_name = f'test_{pinns_method}_{scheme}'
         yaml_path = os.path.join(os.path.dirname(__file__),'simulations_yaml',sim_name+'.yaml')
+        yaml_prev_path = os.path.join(os.path.dirname(__file__),'simulations_yaml','test_born_ion.yaml')
+        shutil.copy(yaml_prev_path,yaml_path)
         results_path = os.path.join(temp_dir,'results',sim_name)
         sim = Simulation(yaml_path, results_path=results_path, molecule_dir=None)
+        sim.pinns_method = pinns_method
+        sim.pbe_model = model
+        sim.equation = scheme
+        if pinns_method == 'DBM':
+            sim.num_networks = 1
+            sim.losses = ['IB1','IB2']
         sim.create_simulation()
         sim.adapt_model()
         sim.solve_model()
@@ -78,30 +90,40 @@ def test_xppbe_solver(molecule):
 
 
 @pytest.mark.parametrize(
- ('loss'),
+ ('molecule'),
  (
-     ('K2'),
-     ('E2'),
-     ('G'),
-     ('Ir')
+     ('methanol'),
+     ('arg')
  )       
 )
-def test_additional_losses(loss):
+def test_other_molecules(molecule):
     with tempfile.TemporaryDirectory() as temp_dir:
-        sim_name = f'test_{loss}'
+        sim_name = f'test_{molecule}'
+        yaml_path = os.path.join(os.path.dirname(__file__),'simulations_yaml',sim_name+'.yaml')
+        results_path = os.path.join(temp_dir,'results',sim_name)
+        sim = Simulation(yaml_path, results_path=results_path, molecule_dir=None)
+        sim.create_simulation()
+        sim.adapt_model()
+        sim.solve_model()
+        sim.postprocessing(run_all=True)
+        run_checkers(sim,sim_name,temp_dir)
+
+
+def test_optimizers():
+    with tempfile.TemporaryDirectory() as temp_dir:
+        sim_name = f'test_optimizers'
         yaml_path = os.path.join(os.path.dirname(__file__),'simulations_yaml',sim_name+'.yaml')
         yaml_prev_path = os.path.join(os.path.dirname(__file__),'simulations_yaml','test_born_ion.yaml')
         shutil.copy(yaml_prev_path,yaml_path)
         results_path = os.path.join(temp_dir,'results',sim_name)
         sim = Simulation(yaml_path, results_path=results_path, molecule_dir=None)
-        sim.losses.append(loss)
-        if loss == 'E2':
-            sim.mesh_properties['dR_exterior'] = 5
+        sim.options_optimizer2['maxiter'] = 3
+        sim.N_steps_2 = 2
         sim.create_simulation()
         sim.adapt_model()
         sim.solve_model()
         sim.postprocessing(run_all=True)
-        run_checkers(sim,sim_name,temp_dir)
+        run_checkers(sim,sim_name,temp_dir,num_iters=2)
 
 
 @pytest.mark.parametrize(
@@ -127,24 +149,26 @@ def test_other_architectures(arch):
         run_checkers(sim,sim_name,temp_dir)
 
 
+
 @pytest.mark.parametrize(
- ('model','scheme'),
+ ('loss'),
  (
-     ('nonlinear','regularized_scheme_2'),
-     ('linear', 'regularized_scheme_1'),
-     ('linear','standard')
+     ('K2'),
+     ('E2'),
+     ('G')
  )       
 )
-def test_non_linear_and_schemes(model,scheme):
+def test_additional_losses(loss):
     with tempfile.TemporaryDirectory() as temp_dir:
-        sim_name = f'test_{model}_{scheme}'
+        sim_name = f'test_{loss}'
         yaml_path = os.path.join(os.path.dirname(__file__),'simulations_yaml',sim_name+'.yaml')
         yaml_prev_path = os.path.join(os.path.dirname(__file__),'simulations_yaml','test_born_ion.yaml')
         shutil.copy(yaml_prev_path,yaml_path)
         results_path = os.path.join(temp_dir,'results',sim_name)
         sim = Simulation(yaml_path, results_path=results_path, molecule_dir=None)
-        sim.pbe_model = model
-        sim.equation = scheme
+        sim.losses.append(loss)
+        if loss == 'E2':
+            sim.mesh_properties['dR_exterior'] = 5
         sim.create_simulation()
         sim.adapt_model()
         sim.solve_model()

diff --git a/xppbe/Mesh/Mesh.py b/xppbe/Mesh/Mesh.py
@@ -95,7 +95,7 @@ class Domain_Mesh():
     DTYPE = 'float32'
     pi = np.pi
 
-    def __init__(self, molecule, mesh_properties, simulation='Main', path='', save_points=False, results_path='', molecule_dir=''):
+    def __init__(self, molecule, mesh_properties, simulation='Main', path='', save_points=False, results_path='', molecule_dir='', losses_names=list()):
 
         self.mesh_properties = {
             'density_mol': 10,
@@ -123,6 +123,7 @@ def __init__(self, molecule, mesh_properties, simulation='Main', path='', save_p
         self.save_points = save_points
         self.main_path = path
         self.simulation_name = simulation
+        self.losses_names = losses_names
         self.results_path = self.main_path if results_path=='' else results_path
         self.molecule_path = os.path.join(self.main_path,'Molecules',molecule) if molecule_dir=='' else molecule_dir
 
@@ -142,10 +143,14 @@ def __init__(self, molecule, mesh_properties, simulation='Main', path='', save_p
 
     def read_create_meshes(self):
         self.create_molecule_mesh()
-        self.create_sphere_mesh()
-        self.create_interior_mesh()
-        self.create_exterior_mesh()
-        self.create_charges_mesh()
+        if 'D2' in self.losses_names or 'R2' in self.losses_names:
+            self.create_sphere_mesh()
+        if 'R1' in self.losses_names:
+            self.create_interior_mesh()
+        if 'R2' in self.losses_names:
+            self.create_exterior_mesh()
+        if 'Q1' in self.losses_names:
+            self.create_charges_mesh()
         self.create_mesh_obj()
         print("Mesh initialization ready")
 
@@ -269,8 +274,11 @@ def create_mesh_obj(self):
         mol_min, mol_max = np.min(self.mol_verts, axis=0), np.max(self.mol_verts, axis=0)
         self.scale_1 = [mol_min.tolist(), mol_max.tolist()]
 
-        sphere_min, sphere_max = np.min(self.sphere_mesh.vertices, axis=0), np.max(self.sphere_mesh.vertices, axis=0)
-        self.scale_2 = [sphere_min.tolist(), sphere_max.tolist()]
+        if 'D2' in self.losses_names or 'R2' in self.losses_names:
+            sphere_min, sphere_max = np.min(self.sphere_mesh.vertices, axis=0), np.max(self.sphere_mesh.vertices, axis=0)
+            self.scale_2 = [sphere_min.tolist(), sphere_max.tolist()]
+        else:
+            self.scale_2 = self.scale_1
 
         if self.save_points:
             X_plot = dict()
@@ -303,13 +311,13 @@ def adapt_meshes_domain(self,data,q_list):
                         X_plot[f'{type_b}_verts'] = X.numpy()
                         self.save_data_plot(X_plot)
 
-            elif type_b in ('Iu','Id','Ir'):
+            elif type_b[0] == 'I':
                 N = self.mol_verts_normal
                 X = tf.constant(self.mol_verts, dtype=self.DTYPE)
                 X_I = (X, N)
                 self.domain_mesh_names.add(type_b)
-                self.domain_mesh_data['I'] = (X_I,flag)
-            
+                self.domain_mesh_data['I'] = (X_I,flag)               
+
             elif type_b in ('G'):
                 self.domain_mesh_names.add(type_b)
                 self.domain_mesh_data[type_b] = ((tf.constant(self.mol_faces_centroid, dtype=self.DTYPE),