add option for mass retention

config.py

@@ -0,0 +1 @@
+water_fraction = False

neural_network.py

@@ -10,6 +10,7 @@ from keras.utils import plot_model
 from matplotlib import pyplot as plt
 
 from CustomScaler import CustomScaler
+from config import water_fraction
 from simulation_list import SimulationList
 
 simulations = SimulationList.jsonlines_load()
@@ -27,7 +28,10 @@ scaler = CustomScaler()
 scaler.fit(X)
 x = scaler.transform_data(X)
 print(x.shape)
-Y = np.array([s.water_retention_both for s in train_data])
+if water_fraction:
+    Y = np.array([s.water_retention_both for s in train_data])
+else:
+    Y = np.array([s.mass_retention_both for s in train_data])
 print(Y.shape)
 X_test = np.array(
     [[s.alpha, s.v, s.projectile_mass, s.gamma, s.target_water_fraction, s.projectile_water_fraction] for s in
@@ -43,9 +47,10 @@ tbCallBack = keras.callbacks.TensorBoard(log_dir='./logs/{}'.format(random.randi
                                          write_grads=True, write_images=True, embeddings_freq=0,
                                          embeddings_layer_names=None, embeddings_metadata=None, embeddings_data=None,
                                          update_freq='epoch')
+modelname = "model.hd5" if water_fraction else "model_mass.hd5"
 
-if os.path.exists("model.hd5"):
-    model = load_model("model.hd5")
+if os.path.exists(modelname):
+    model = load_model(modelname)
 else:
     model = Sequential()
     model.add(Dense(6, input_dim=6, activation='relu'))
@@ -64,16 +69,16 @@ else:
         ...
     # print("-------------------------------------")
     # exit()
-    model.save("model.hd5")
+    model.save(modelname)
 
-xrange = np.linspace(-0.5, 60.5, 100)
-yrange = np.linspace(0.5, 5.5, 100)
+xrange = np.linspace(-0.5, 60.5, 300)
+yrange = np.linspace(0.5, 5.5, 300)
 xgrid, ygrid = np.meshgrid(xrange, yrange)
 mcode = 1e24
 wpcode = 15 / 100
 wtcode = 15 / 100
 gammacode = 0.6
-testinput = np.array([[np.nan, np.nan, mcode, gammacode, wtcode, wpcode]] * 100 * 100)
+testinput = np.array([[np.nan, np.nan, mcode, gammacode, wtcode, wpcode]] * 300 * 300)
 testinput[::, 0] = xgrid.flatten()
 testinput[::, 1] = ygrid.flatten()
 testinput = scaler.transform_data(testinput)
@@ -81,11 +86,11 @@ testinput = scaler.transform_data(testinput)
 print(testinput)
 print(testinput.shape)
 testoutput = model.predict(testinput)
-outgrid = np.reshape(testoutput, (100, 100))
+outgrid = np.reshape(testoutput, (300, 300))
 print("minmax")
 print(np.nanmin(outgrid), np.nanmax(outgrid))
-
-plt.imshow(outgrid, interpolation='none', cmap="Blues", aspect="auto", origin="lower", vmin=0, vmax=1,
+cmap = "Blues" if water_fraction else "Oranges"
+plt.imshow(outgrid, interpolation='none', cmap=cmap, aspect="auto", origin="lower", vmin=0, vmax=1,
            extent=[xgrid.min(), xgrid.max(), ygrid.min(), ygrid.max()])
 
 plt.colorbar().set_label("water retention fraction")

simulation.py

@@ -59,10 +59,22 @@ class Simulation:
         """
         return self.second_largest_aggregate_mass / self.projectile_mass
 
+    @property
+    def projectile_stone_fraction(self):
+        return 1 - self.projectile_water_fraction
+
+    @property
+    def target_stone_fraction(self):
+        return 1 - self.target_water_fraction
+
     @property
     def initial_water_mass(self) -> float:
         return self.projectile_mass * self.projectile_water_fraction + self.target_mass * self.target_water_fraction
 
+    @property
+    def initial_stone_mass(self) -> float:
+        return self.projectile_mass * self.projectile_stone_fraction + self.target_mass * self.target_stone_fraction
+
     @property
     def water_retention_both(self) -> float:
         """
@@ -75,7 +87,10 @@ class Simulation:
 
     @property
     def mass_retention_both(self) -> float:
-        return (self.largest_aggregate_mass + self.second_largest_aggregate_mass) / self.total_mass
+        return (
+                       self.largest_aggregate_mass * (1 - self.largest_aggregate_water_fraction)
+                       + self.second_largest_aggregate_mass * (1 - self.second_largest_aggregate_water_fraction)
+               ) / (self.projectile_mass * (1-self.projectile_water_fraction) + self.target_mass * (1-self.target_water_fraction))
 
     @property
     def water_retention_main(self) -> float:

simulation_list.py

@@ -4,6 +4,7 @@ from typing import List
 
 import numpy as np
 
+from config import water_fraction
 from simulation import Simulation
 
 
@@ -62,4 +63,12 @@ class SimulationList:
 
     @property
     def Y(self):
+        return self.Y_water if water_fraction else self.Y_mass
+
+    @property
+    def Y_mass(self):
+        return np.array([s.mass_retention_both for s in self.simlist if not s.testcase])
+
+    @property
+    def Y_water(self):
         return np.array([s.water_retention_both for s in self.simlist if not s.testcase])

testing.py

@@ -5,6 +5,7 @@ import numpy as np
 from keras.engine.saving import load_model
 
 from CustomScaler import CustomScaler
+from config import water_fraction
 from interpolators.griddata import GriddataInterpolator
 from interpolators.rbf import RbfInterpolator
 from simulation import Simulation
@@ -14,7 +15,7 @@ simulations = SimulationList.jsonlines_load()
 
 scaler = CustomScaler()
 scaler.fit(simulations.X)
-model = load_model("model.hd5")
+model = load_model("model.hd5" if water_fraction else "model_mass.hd5")
 
 
 def squared_error(inter: float, correct: float) -> float:
@@ -53,8 +54,9 @@ rbf_squared_errors = []
 rbf_errors = []
 grid_squared_errors = []
 grid_errors = []
+cachefile="grid-testing-cache.json" if water_fraction else "grid-testing-cache-mass.json"
 try:
-    with open("grid-testing-cache.json") as f:
+    with open(cachefile) as f:
         raw_data = json.load(f)
         grid_testing_cache = {int(key): value for key, value in raw_data.items()}
 except FileNotFoundError:
@@ -83,7 +85,7 @@ for sim in simulations.simlist:
         if np.isnan(grid_output):
             grid_output = False
         grid_testing_cache[sim.runid] = grid_output
-        with open("grid-testing-cache.json", "w") as f:
+        with open(cachefile, "w") as f:
             json.dump(grid_testing_cache, f)
     if grid_output:
         grid_squared_errors.append(squared_error(grid_output, sim.water_retention_both))

visualize.py

@@ -2,12 +2,14 @@ import numpy as np
 from matplotlib import pyplot as plt, cm
 
 from CustomScaler import CustomScaler
+from config import water_fraction
 from interpolators.griddata import GriddataInterpolator
+from interpolators.rbf import RbfInterpolator
 from simulation_list import SimulationList
 
 
 def main():
-    mcode, gamma, wt, wp = [10 ** 22, 0.6, 15 / 100, 15 / 100]
+    mcode, gamma, wt, wp = [10 ** 24, 0.6, 15 / 100, 15 / 100]
     simlist = SimulationList.jsonlines_load()
     # for s in simlist.simlist:
     #     if s.type!="original":
@@ -39,7 +41,7 @@ def main():
     print(np.nanmin(grid_result), np.nanmax(grid_result))
 
     # plt.title("m={:3.0e}, gamma={:3.1f}, wt={:2.0f}%, wp={:2.0f}%\n".format(mcode, gamma, wt*100, wp*100))
-    cmap = cm.Blues
+    cmap = cm.Blues if water_fraction else cm.Oranges
     cmap.set_bad('white', 1.)  # show nan white
     # plt.contourf(grid_alpha, grid_v, grid_result, 100, cmap="Blues", vmin=0, vmax=1)
     # plt.pcolormesh(grid_alpha, grid_v, grid_result, cmap="Blues", vmin=0, vmax=1)
@@ -51,7 +53,7 @@ def main():
     plt.ylabel("velocity $v$ [$v_{esc}$]")
     plt.tight_layout()
     # plt.savefig("vis.png", transparent=True)
-    plt.savefig("../arbeit/images/plots/griddata1.pdf")
+    plt.savefig("../arbeit/images/plots/mass_griddata2.pdf")
     plt.show()