diff --git a/src/py/mat3ra/made/tools/analyze.py b/src/py/mat3ra/made/tools/analyze.py
index afb3c41c..15f49ad6 100644
--- a/src/py/mat3ra/made/tools/analyze.py
+++ b/src/py/mat3ra/made/tools/analyze.py
@@ -417,7 +417,6 @@ def get_coordination_numbers(
     return coordination_numbers
 
 
-# TODO: fix decorator removal of added atoms and uncomment
 @decorator_handle_periodic_boundary_conditions(cutoff=0.1)
 def get_undercoordinated_atom_indices(
     material: Material,
diff --git a/src/py/mat3ra/made/tools/utils/__init__.py b/src/py/mat3ra/made/tools/utils/__init__.py
index 392ee17b..4c3dbd0b 100644
--- a/src/py/mat3ra/made/tools/utils/__init__.py
+++ b/src/py/mat3ra/made/tools/utils/__init__.py
@@ -166,19 +166,18 @@ def augment_material_with_periodic_images(material: Material, cutoff: float = 0.
     original_count = len(material.basis.coordinates.values)
     coordinates = np.array(material.basis.coordinates.values)
     elements = np.array(material.basis.elements.values)
-    augmented_coords = []
-    augmented_elems = []
+    augmented_coords = coordinates.tolist()
+    augmented_elems = elements.tolist()
 
     for axis in range(3):
         for direction in [-1, 1]:
             translated_coords, translated_elems = filter_and_translate(coordinates, elements, axis, cutoff, direction)
-            augmented_coords.append(translated_coords)
-            augmented_elems.append(translated_elems)
+            augmented_coords.extend(translated_coords)
+            augmented_elems.extend(translated_elems)
+            coordinates = np.array(augmented_coords)
+            elements = np.array(augmented_elems)
 
-    augmented_coords = np.vstack(augmented_coords).tolist()
-    augmented_elems = np.concatenate(augmented_elems).tolist()
     augmented_material = material.clone()
-
     new_basis = augmented_material.basis.copy()
     for i, coord in enumerate(augmented_coords):
         new_basis.add_atom(augmented_elems[i], coord)