update to test in servers

problems/function_fitting/__init__.py

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

problems/function_fitting/diy.py

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

problems/function_fitting/function_fitting_problem.py

@@ -15,8 +15,25 @@ class FunctionFittingProblem(Problem):
         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))
+    def evaluate(self, pop_batch_forward):
+        outs = pop_batch_forward(self.inputs)
+        outs = jax.device_get(outs)
+        fitnesses = -np.mean((self.target - outs) ** 2, axis=(1, 2))
         return fitnesses.tolist()
+
+    def draw(self, batch_func):
+        outs = batch_func(self.inputs)
+        outs = jax.device_get(outs)
+        print(outs)
+        from matplotlib import pyplot as plt
+        plt.xlabel('x')
+        plt.ylabel('y')
+        plt.plot(self.inputs, self.target, color='red', label='target')
+        plt.plot(self.inputs, outs, color='blue', label='predict')
+        plt.legend()
+        plt.show()
+
+    def print(self, batch_func):
+        outs = batch_func(self.inputs)
+        outs = jax.device_get(outs)
+        print(outs)

problems/function_fitting/sin.py

@@ -8,7 +8,7 @@ class Sin(FunctionFittingProblem):
         self.num_inputs = 1
         self.num_outputs = 1
         self.batch = size
-        self.inputs = np.linspace(0, np.pi, self.batch)[:, None]
+        self.inputs = np.linspace(0, 2 * 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)