add package problems

problems/function_fitting/__init__.py

@@ -0,0 +1,3 @@
+from .function_fitting_problem import FunctionFittingProblem
+from .xor import *
+from .sin import *

problems/function_fitting/function_fitting_problem.py

@@ -0,0 +1,22 @@
+import numpy as np
+import jax
+
+from problems import Problem
+
+
+class FunctionFittingProblem(Problem):
+    def __init__(self, num_inputs, num_outputs, batch, inputs, target, loss='MSE'):
+        self.forward_way = 'pop_batch'
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+        self.batch = batch
+        self.inputs = inputs
+        self.target = target
+        self.loss = loss
+        super().__init__(self.forward_way, self.num_inputs, self.num_outputs, self.batch)
+
+    def evaluate(self, batch_forward_func):
+        out = batch_forward_func(self.inputs)
+        out = jax.device_get(out)
+        fitnesses = 1 - np.mean((self.target - out) ** 2, axis=(1, 2))
+        return fitnesses.tolist()

problems/function_fitting/sin.py

@@ -0,0 +1,14 @@
+import numpy as np
+
+from . import FunctionFittingProblem
+
+
+class Sin(FunctionFittingProblem):
+    def __init__(self, size=100):
+        self.num_inputs = 1
+        self.num_outputs = 1
+        self.batch = size
+        self.inputs = np.linspace(0, np.pi, self.batch)[:, None]
+        self.target = np.sin(self.inputs)
+        print(self.inputs, self.target)
+        super().__init__(self.num_inputs, self.num_outputs, self.batch, self.inputs, self.target)

problems/function_fitting/xor.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+from . import FunctionFittingProblem
+
+
+class Xor(FunctionFittingProblem):
+    def __init__(self):
+        self.num_inputs = 2
+        self.num_outputs = 1
+        self.batch = 4
+        self.inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32)
+        self.target = np.array([[0], [1], [1], [0]], dtype=np.float32)
+        super().__init__(self.num_inputs, self.num_outputs, self.batch, self.inputs, self.target)