Revert unworkable structure

scico/flax/examples/data_generation.py

@@ -12,26 +12,23 @@
 
 from functools import partial
 from time import time
-from typing import Callable, Tuple
+from typing import Callable, List, Tuple, Union
 
 import numpy as np
 
 try:
     import xdesign  # noqa: F401
 except ImportError:
     have_xdesign = False
-    # pylint: disable=missing-function-docstring
 
-    def generate_foam1_images():
-        raise RunTimeError("xdesign package required.")
+    # pylint: disable=missing-class-docstring
+    class UnitCircle:
+        pass
 
-    def generate_foam2_images():
-        raise RunTimeError("xdesign package required.")
-
-    # pylint: enable=missing-function-docstring
+    # pylint: enable=missing-class-docstring
 else:
     have_xdesign = True
-    from .xdesign_func import generate_foam1_images, generate_foam2_images
+    from xdesign import Foam, SimpleMaterial, UnitCircle, discrete_phantom
 
 try:
     import ray  # noqa: F401
@@ -55,6 +52,96 @@ def generate_foam2_images():
     from scico.linop.xray.astra import XRayTransform2D
 
 
+class Foam2(UnitCircle):
+    """Foam-like material with two attenuations.
+
+    Define functionality to generate phantom with structure similar
+    to foam with two different attenuation properties."""
+
+    def __init__(
+        self,
+        size_range: Union[float, List[float]] = [0.05, 0.01],
+        gap: float = 0,
+        porosity: float = 1,
+        attn1: float = 1.0,
+        attn2: float = 10.0,
+    ):
+        """Foam-like structure with two different attenuations.
+        Circles for material 1 are more sparse than for material 2
+        by design.
+
+        Args:
+            size_range: The radius, or range of radius, of the
+                circles to be added. Default: [0.05, 0.01].
+            gap: Minimum distance between circle boundaries.
+                Default: 0.
+            porosity: Target porosity. Must be a value between
+                [0, 1]. Default: 1.
+            attn1: Mass attenuation parameter for material 1.
+                Default: 1.
+            attn2: Mass attenuation parameter for material 2.
+                Default: 10.
+        """
+        super().__init__(radius=0.5, material=SimpleMaterial(attn1))
+        if porosity < 0 or porosity > 1:
+            raise ValueError("Porosity must be in the range [0,1).")
+        self.sprinkle(
+            300, size_range, gap, material=SimpleMaterial(attn2), max_density=porosity / 2.0
+        ) + self.sprinkle(300, size_range, gap, material=SimpleMaterial(20), max_density=porosity)
+
+
+def generate_foam1_images(seed: float, size: int, ndata: int) -> np.ndarray:
+    """Generate batch of xdesign foam-like structures.
+
+    Generate batch of images with `xdesign` foam-like structure, which
+    uses one attenuation.
+
+    Args:
+        seed: Seed for data generation.
+        size: Size of image to generate.
+        ndata: Number of images to generate.
+
+    Returns:
+        Array of generated data.
+    """
+    if not have_xdesign:
+        raise RuntimeError("Package xdesign is required for use of this module.")
+    np.random.seed(seed)
+    saux = np.zeros((ndata, size, size, 1), dtype=np.float32)
+    for i in range(ndata):
+        foam = Foam(size_range=[0.075, 0.0025], gap=1e-3, porosity=1)
+        saux[i, ..., 0] = discrete_phantom(foam, size=size)
+
+    return saux
+
+
+def generate_foam2_images(seed: float, size: int, ndata: int) -> np.ndarray:
+    """Generate batch of foam2 structures.
+
+    Generate batch of images with :class:`Foam2` structure
+    (foam-like material with two different attenuations).
+
+    Args:
+        seed: Seed for data generation.
+        size: Size of image to generate.
+        ndata: Number of images to generate.
+
+    Returns:
+        Array of generated data.
+    """
+    if not have_xdesign:
+        raise RuntimeError("Package xdesign is required for use of this module.")
+    np.random.seed(seed)
+    saux = np.zeros((ndata, size, size, 1), dtype=np.float32)
+    for i in range(ndata):
+        foam = Foam2(size_range=[0.075, 0.0025], gap=1e-3, porosity=1)
+        saux[i, ..., 0] = discrete_phantom(foam, size=size)
+    # normalize
+    saux /= np.max(saux, axis=(1, 2), keepdims=True)
+
+    return saux
+
+
 def vector_f(f_: Callable, v: Array) -> Array:
     """Vectorize application of operator.