Merge remote-tracking branch 'origin/main'

Author: Leguark

gempy/core/color_generator.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from functools import cached_property
 from typing import Optional
 
 import numpy as np
@@ -36,11 +37,10 @@ def __init__(self):
     @staticmethod
     def _random_hexcolor() -> str:
         """Generates a random hex color string."""
-        return "#"+str(hex(np.random.randint(0, 16777215))).lstrip("0x")
+        return f"#{np.random.randint(0, 0xffffff):06x}"
 
     def regenerate_color_palette(self, seaborn_palettes: Optional[list[str]] = None):
         try:
-            
             import seaborn as sns
             seaborn_installed = True
         except ImportError:
@@ -54,7 +54,7 @@ def regenerate_color_palette(self, seaborn_palettes: Optional[list[str]] = None)
         elif seaborn_palettes and not seaborn_installed:
             raise ImportError("Seaborn is not installed. Please install it to use color palettes.")
         else:
-            hex_colors = self._gempy_default_colors
+            hex_colors = list(self._gempy_default_colors)
 
         self.hex_colors = hex_colors
 
@@ -65,14 +65,18 @@ def __iter__(self) -> 'ColorsGenerator':
 
     def __next__(self) -> str:
         """Generator that yields the next color."""
-        for color in self.hex_colors:
-            result = self.hex_colors[self._index]
-            self._index += 1
-            return result
+        color = self.up_next()
+        self._index += 1
+        del self._next_color
+        return color
 
-        while True:
-            return self._random_hexcolor()
-    
     def up_next(self) -> str:
         """Returns the next color without incrementing the index."""
-        return self.hex_colors[self._index]
+        return self._next_color
+
+    @cached_property
+    def _next_color(self) -> str:
+        if self._index < len(self.hex_colors):
+            return self.hex_colors[self._index]
+
+        return self._random_hexcolor()

gempy/modules/data_manipulation/manipulate_points.py

@@ -94,8 +94,8 @@ def add_orientations(geo_model: GeoModel,
         y: Sequence[float],
         z: Sequence[float],
         elements_names: Sequence[str],
-        pole_vector: Optional[Sequence[np.ndarray]] = None,
-        orientation: Optional[Sequence[np.ndarray]] = None,
+        pole_vector: Optional[Union[Sequence[np.ndarray], np.ndarray]] = None,
+        orientation: Optional[Union[Sequence[np.ndarray], np.ndarray]] = None,
         nugget: Optional[Sequence[float]] = None
 ) -> StructuralFrame:
     """Add orientation data to the geological model.
@@ -110,8 +110,8 @@ def add_orientations(geo_model: GeoModel,
         y (Sequence[float]): Sequence of y-coordinates for the orientation points.
         z (Sequence[float]): Sequence of z-coordinates for the orientation points.
         elements_names (Sequence[str]): Sequence of element names corresponding to each orientation point.
-        pole_vector (Optional[Sequence[np.ndarray]]): Sequence of pole vectors for the orientation points.
-        orientation (Optional[Sequence[np.ndarray]]): Sequence of orientation angles (azimuth, dip, polarity) for the orientation points.
+        pole_vector (Optional[Union[Sequence[np.ndarray], np.ndarray]]): Sequence of pole vectors for each orientation point. If is np.ndarray, it should have shape (n, 3).
+        orientation (Optional[Union[Sequence[np.ndarray], np.ndarray]]): Sequence of orientation angles for each orientation point. If is np.ndarray, it should have shape (n, 3).
         nugget (Optional[Sequence[float]]): Sequence of nugget values for each orientation point. If not provided,
             a default value will be used for all points.
 
@@ -125,7 +125,8 @@ def add_orientations(geo_model: GeoModel,
     if pole_vector is None and orientation is None:
         raise ValueError("Either pole_vector or orientation must be provided.")
 
-    if orientation:  # Convert orientation to pole_vector (or gradient)
+    if orientation is not None:
+        orientation = np.array(orientation)
         pole_vector = convert_orientation_to_pole_vector(
             azimuth=orientation[:, 0],
             dip=orientation[:, 1],

test/test_core/test_color_generator.py

@@ -1,10 +1,10 @@
-from gempy.core.color_generator import ColorsGenerator
+import gempy.core.color_generator as cg
 
 
 def test_color_generator():
 
     # create an instance of ColorsGenerator
-    color_iterator: iter = ColorsGenerator()
+    color_iterator: iter = cg.ColorsGenerator()
 
     # generate color dictionary (with optional seaborn palettes if you want)
 
@@ -13,4 +13,32 @@ def test_color_generator():
     # now you can get colors one by one
     print(next(color_iterator))  # prints first color
     print(next(color_iterator))  # prints second color
+
+    left_colors = color_iterator._gempy_default_colors[2:]
+
+    assert left_colors == [
+        color for color, _ in zip(color_iterator, left_colors)
+    ]
+
+    # get a random color
+    assert color_iterator.up_next() == next(color_iterator)
+
+    # override the random color generator
+    cg.np.random.randint = lambda a, b: 0x00ff42
+
+    assert color_iterator._random_hexcolor() == "#00ff42"
+    assert next(color_iterator) == "#00ff42"
+
+    # trigger up_next method for the cache
+    assert color_iterator.up_next() == "#00ff42"
+
+    cg.np.random.randint = lambda a, b: 0x000000
+
+    # ensure that the up_next method still returns the same color
+    assert color_iterator.up_next() == "#00ff42"
+
+    # trigger the cache
+    assert next(color_iterator) == "#00ff42"
+    assert next(color_iterator) == "#000000"
+
     # etc...