From eaefc69be0ec9f16b1bcd70f31bc2fc74949e2a4 Mon Sep 17 00:00:00 2001
From: Lukas Winkler <git@lw1.at>
Date: Wed, 9 Jun 2021 15:39:05 +0200
Subject: [PATCH] add neural network base massloss method

---
 massloss/__init__.py           |   1 +
 massloss/simple_nn_massloss.py | 105 +++++++++++++++++++++++++++++++++
 merge.py                       |   4 +-
 3 files changed, 108 insertions(+), 2 deletions(-)
 create mode 100644 massloss/simple_nn_massloss.py

diff --git a/massloss/__init__.py b/massloss/__init__.py
index 8f6cc02..1d4286a 100644
--- a/massloss/__init__.py
+++ b/massloss/__init__.py
@@ -1,3 +1,4 @@
 from .base_massloss import *
 from .lei_zhou_massloss import *
 from .rbf_massloss import *
+from .simple_nn_massloss import *
diff --git a/massloss/simple_nn_massloss.py b/massloss/simple_nn_massloss.py
new file mode 100644
index 0000000..c2a2710
--- /dev/null
+++ b/massloss/simple_nn_massloss.py
@@ -0,0 +1,105 @@
+import json
+from dataclasses import dataclass
+from math import exp
+from typing import Tuple, List
+
+from massloss import Massloss
+
+Layer = List[float]
+
+
+def relu(x):
+    return max(0, x)
+
+
+def sigmoid(x):
+    return 1 / (1 + exp(-x))
+
+
+@dataclass
+class Model:
+    """
+    see https://github.com/Findus23/nn_evaluate
+    for more implementations of this model in other languages
+    """
+    means: List[float]  # 6
+    stds: List[float]  # 6
+    hidden_weight: List[List[float]]  # 50x6
+    hidden_bias: List[float]  # 50
+    output_weight: List[List[float]]  # 3x50
+    output_bias: List[float]  # 3
+
+    @property
+    def hidden_layer_size(self):
+        return len(self.hidden_bias)
+
+    @property
+    def input_layer_size(self):
+        return len(self.means)
+
+    @property
+    def output_layer_size(self):
+        return len(self.output_bias)
+
+    def calculate_layer(self, layer_size, parent_layer_size, parent_layer, weight, bias) -> Layer:
+        new_layer = []
+        for hl in range(layer_size):
+            node = 0
+            for parent in range(parent_layer_size):
+                node += parent_layer[parent] * weight[hl][parent]
+            node += bias[hl]
+            new_layer.append(node)
+        return new_layer
+
+    def scale_input(self, input: List[float]):
+        result = []
+        for index, parameter in enumerate(input):
+            result.append((parameter - self.means[index]) / self.stds[index])
+        return result
+
+    def evaluate(self, input: List[float]):
+        scaled_input = self.scale_input(input)
+        hidden_layer = self.calculate_layer(
+            layer_size=self.hidden_layer_size,
+            parent_layer_size=self.input_layer_size,
+            parent_layer=scaled_input,
+            weight=self.hidden_weight,
+            bias=self.hidden_bias
+        )
+        hidden_layer = [relu(x) for x in hidden_layer]
+        output_layer = self.calculate_layer(
+            layer_size=self.output_layer_size,
+            parent_layer_size=self.hidden_layer_size,
+            parent_layer=hidden_layer,
+            weight=self.output_weight,
+            bias=self.output_bias
+        )
+
+        output_layer = [sigmoid(x) for x in output_layer]
+        return output_layer
+
+
+class SimpleNNMassloss(Massloss):
+    name = "simpleNN"
+
+    def __init__(self):
+        with open("./pytorch_model.json") as f:
+            data = json.load(f)
+        self.model = Model(
+            hidden_weight=data["hidden.weight"],
+            hidden_bias=data["hidden.bias"],
+            output_weight=data["output.weight"],
+            output_bias=data["output.bias"],
+            means=data["means"],
+            stds=data["stds"],
+        )
+        self.wp = self.wt = 1e-4
+
+    def estimate(self, alpha, velocity, projectile_mass, gamma) -> Tuple[float, float, float]:
+        result = self.model.evaluate([alpha, velocity, projectile_mass, gamma, self.wt, self.wp])
+        return result[0], result[1], result[2]
+
+
+if __name__ == '__main__':
+    inter = SimpleNNMassloss()
+    print(inter.estimate(32, 3, 7.6e22, 0.16))
diff --git a/merge.py b/merge.py
index a7c9d6f..e0da492 100644
--- a/merge.py
+++ b/merge.py
@@ -9,7 +9,7 @@ from rebound.simulation import POINTER_REB_SIM, reb_collision
 from scipy.constants import astronomical_unit, G
 
 from extradata import ExtraData, ParticleData, CollisionMeta, Input
-from massloss import RbfMassloss, Massloss, LeiZhouMassloss
+from massloss import RbfMassloss, Massloss, LeiZhouMassloss, SimpleNNMassloss
 from massloss.perfect_merging import PerfectMerging
 from utils import unique_hash, clamp, PlanetaryRadius
 
@@ -126,7 +126,7 @@ def merge_particles(sim_p: POINTER_REB_SIM, collision: reb_collision, ed: ExtraD
     print("interpolating")
 
     if not massloss_estimator:
-        methods = [RbfMassloss, LeiZhouMassloss, PerfectMerging]
+        methods = [RbfMassloss, LeiZhouMassloss, PerfectMerging, SimpleNNMassloss]
         per_name = {}
         for method in methods:
             per_name[method.name] = method