From deb42f64ad524d6a3758c6a6e2062382fafefeb6 Mon Sep 17 00:00:00 2001
From: alperenkaran <62502319+alperenkaran@users.noreply.github.com>
Date: Sun, 12 Jun 2022 17:08:41 +0300
Subject: [PATCH] added eyeglasses shape

---
 README.md           |  2 ++
 tadasets/shapes.py  | 78 ++++++++++++++++++++++++++++++++++++++++++++-
 test/test_shapes.py | 19 +++++++++++
 3 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 277e7c4..b2eb4b5 100644
--- a/README.md
+++ b/README.md
@@ -24,6 +24,7 @@ The shape constructors are exposed in a functional interface, each returning a n
 - d-sphere
 - swiss roll
 - infinity sign
+- eyeglasses
 
 Each shape can be embedded in arbitrary ambient dimension by supplying the `ambient` argument. Additionally, noise can be added to the shape through the `noise` argument.
 
@@ -34,6 +35,7 @@ torus = tadasets.torus(n=2000, c=2, a=1, ambient=200, noise=0.2)
 swiss_roll = tadasets.swiss_roll(n=2000, r=4, ambient=10, noise=1.2)
 dsphere = tadasets.dsphere(n=1000, d=12, r=3.14, ambient=14, noise=0.14)
 infty_sign = tadasets.infty_sign(n=3000, noise=0.1)
+eyeglasses = tadasets.eyeglasses(n=1000, r1=1, r2=2, neck_size=.5, noise=0.1, ambient=10)
 ```
 
 Contributions
diff --git a/tadasets/shapes.py b/tadasets/shapes.py
index 28516ee..d7378df 100644
--- a/tadasets/shapes.py
+++ b/tadasets/shapes.py
@@ -2,7 +2,7 @@
 from .dimension import embed
 from .rotate import rotate_2D
 
-__all__ = ["torus", "dsphere", "sphere", "swiss_roll", "infty_sign"]
+__all__ = ["torus", "dsphere", "sphere", "swiss_roll", "infty_sign", "eyeglasses"]
 
 
 # TODO: Make a base class that controls `ambient` and `noise`.
@@ -186,3 +186,79 @@ def infty_sign(n=100, noise=None, angle=None, seed=None):
 
         X = rotate_2D(X, angle=angle)
     return X
+
+
+def eyeglasses(n=100, r1=1, r2=None, neck_size=None, noise=None, ambient=None, seed=None):
+    """Sample `n` points on an eyeglasses shape.
+
+    Parameters
+    ----------
+    n : int, default=100
+        Number of points in shape
+    r1 : float, default=1
+        The radius of the left half of the eyeglasses shape
+    r2 : float, default=None
+        The radius of the right half. If None, it is equal to `r1`.
+    neck_size : float, default=None
+        The width of the neck.
+    noise : float, default=None
+        Standard deviation of normally distributed noise added to data.
+    ambient : int, default=None
+        Embed the eyeglasses shape into a space with ambient dimension equal to `ambient`.
+        The eyeglasses shape is randomly rotated in this high dimensional space.
+    seed : int
+        Seed for random state.
+    """
+    np.random.seed(seed)
+
+    if r2 is None:
+        r2 = r1
+    if neck_size is None:
+        neck_size = min(r1, r2)
+
+    assert neck_size < 2 * min(r1, r2), 'Neck should be smaller than 2*min(r1,r2).'
+    half_neck = neck_size / 2
+
+    r_neck = min(r1, r2)
+
+    x_left = ((r_neck + r1) ** 2 - (r_neck + half_neck) ** 2) ** (1 / 2) - r1  # x distance from left circle to neck
+    x_right = ((r_neck + r2) ** 2 - (r_neck + half_neck) ** 2) ** (1 / 2) - r2
+
+    alpha_1 = np.arcsin((r_neck + half_neck) / (r1 + r_neck))
+    alpha_2 = np.arcsin((r_neck + half_neck) / (r2 + r_neck))
+
+    centers = {'l': [-r1 - x_left, 0],
+               'r': [x_right + r2, 0],
+               't': [0, half_neck + r_neck],
+               'b': [0, -half_neck - r_neck]}
+
+    radii = {'l': r1, 'r': r2, 't': r_neck, 'b': r_neck}
+
+    angle_ranges = {'l': [alpha_1, 2 * np.pi - alpha_1],
+                    'r': [-np.pi + alpha_2, np.pi - alpha_2],
+                    't': [np.pi + alpha_1, 2 * np.pi - alpha_2],
+                    'b': [alpha_2, np.pi - alpha_1]}
+
+    # compute how many points each circle will contain.
+    arc_lens = {x: (angle_ranges[x][1] - angle_ranges[x][0]) * radii[x] for x in ['l', 'r', 't', 'b']}
+    total_arc_len = sum(arc_lens.values())
+    probabilities = {x: arc_lens[x] / total_arc_len for x in ['l', 'r', 't', 'b']}
+    positions = np.random.choice(['l', 'r', 't', 'b'], p=list(probabilities.values()), size=n)
+    positions, counts = np.unique(positions, return_counts=True)
+    counts = {x: c for x, c in zip(positions, counts)}
+
+    data = []
+    for x in ['l', 'r', 't', 'b']:
+        angles = np.random.uniform(size=counts[x], low=angle_ranges[x][0], high=angle_ranges[x][1])
+        points = np.vstack((radii[x] * np.cos(angles), radii[x] * np.sin(angles))).T + centers[x]
+        data.append(points)
+
+    data = np.concatenate(data)
+
+    if noise:
+        data += noise * np.random.randn(n, 2)
+
+    if ambient:
+        data = embed(data, ambient)
+
+    return data
diff --git a/test/test_shapes.py b/test/test_shapes.py
index cea6061..6ab985e 100644
--- a/test/test_shapes.py
+++ b/test/test_shapes.py
@@ -115,3 +115,22 @@ def test_rotation(self):
             assert ae is not None
         t = tadasets.infty_sign(n=345, angle=2)
         assert t.shape[0] == 345
+
+
+class TestEyeglasses:
+    def test_n(self):
+        t = tadasets.eyeglasses(n=345, r1=1, r2=2, neck_size=.8)
+        assert t.shape[0] == 345
+
+    def test_neck(self):
+        t = tadasets.eyeglasses(n=5000, r1=1, r2=2, neck_size=.8)
+        top, bottom = t[t[:, 1] > 0], t[t[:, 1] < 0]
+        y_neck_top = top[np.abs(top[:, 0]).argmin(), 1]
+        y_neck_bottom = bottom[np.abs(bottom[:, 0]).argmin(), 1]
+        assert np.abs(y_neck_top - y_neck_bottom - .8) <= .001
+
+    def test_r(self):
+        t = tadasets.eyeglasses(n=5000, r1=1, r2=2, neck_size=.8)
+        left, right = t[t[:, 0] < 0], t[t[:, 0] > 0]
+        assert np.abs(left[:, 1].max() - 1) <= .001
+        assert np.abs(right[:, 1].max() - 2) <= .001