Simplify code to update space groups based on subset

alexhernandezgarcia · Oct 18, 2023 · 75aa91e · 75aa91e
1 parent e2e60dc
commit 75aa91e
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/gflownet/envs/crystals/spacegroup.py b/gflownet/envs/crystals/spacegroup.py
@@ -687,10 +687,7 @@ def _restrict_space_groups(self, sg_subset: Optional[Iterable] = None):
         sg_subset = set(sg_subset)
 
         # Update self.space_groups
-        self.space_groups = deepcopy(self.space_groups)
-        sg_to_remove = [sg for sg in self.space_groups if sg not in sg_subset]
-        for sg in sg_to_remove:
-            del self.space_groups[sg]
+        self.space_groups = [sg for sg in self.space_groups if sg in sg_subset]
 
         # Update self.crystal_lattice_systems based on space groups
         self.crystal_lattice_systems = deepcopy(self.crystal_lattice_systems)