fix type comparison

Helmholtz-AI-Energy · Aug 14, 2024 · 22a4d37 · 22a4d37
1 parent b1b967a
commit 22a4d37
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/propulate/population.py b/propulate/population.py
@@ -99,11 +99,11 @@ def __getitem__(self, key: str) -> Union[float, int, str]:
             return self.mapping[key]
         else:
             # continuous variable
-            if self.types[key] == float:
+            if self.types[key] is float:
                 return float(self.position[self.offsets[key]].item())
-            elif self.types[key] == int:
+            elif self.types[key] is int:
                 return int(np.rint(self.position[self.offsets[key]]).item())
-            elif self.types[key] == str:
+            elif self.types[key] is str:
                 offset = self.offsets[key]
                 upper = self.offsets[key] + len(self.limits[key])
                 return str(
@@ -120,13 +120,13 @@ def __setitem__(self, key: str, newvalue: Union[float, int, str, Any]) -> None:
         else:
             if key not in self.limits:
                 raise ValueError("Unknown gene.")
-            if self.types[key] == float:
+            if self.types[key] is float:
                 assert isinstance(newvalue, float)
                 self.position[self.offsets[key]] = newvalue
-            elif self.types[key] == int:
+            elif self.types[key] is int:
                 assert isinstance(newvalue, int)
                 self.position[self.offsets[key]] = float(newvalue)
-            elif self.types[key] == str:
+            elif self.types[key] is str:
                 assert newvalue in self.limits[key]
                 offset = self.offsets[key]
                 upper = len(self.limits[key])