3D Objects

…

from typing import Optional

import plotly.graph_objects as go
import numpy as np

from examples.util import show

Surface Object

class Surface:
    """
    A surface represented by vertices and triangles (faces) between these vertices
    """

    def __init__(self, vertices=(), faces=()):
        """
        :param vertices: N initial vertices as 3D coordinates; floats of shape (N, 3)
        :param faces: M initial faces as i, j, k indices of vertices; ints of shape (M, 3)
        """
        self.vertices = np.array(vertices)
        self.faces = np.array(faces)

    def refine(self, adjust=None, repeat: Optional[int] = 1):
        """
        Refine the surface by splitting each face into four by inserting new vertices at the edges' midpoints:

           /\              /\
          /  \            /__\
         /    \   ==>    /\  /\
        /______\        /__\/__\

        :param adjust: gets a new vertex as (x, y, z) coordinates initialised as midpoint between existing vertices
         and returns a refined vertex (e.g. projected to the exact surface); in derived classes, if adjust is None and
         the class defines an adjust function, this will be used instead
        :param repeat: number of refinement steps (each steps increases the number of faces by a factor of four)
        """
        # repeat
        if repeat is not None:
            for _ in range(repeat):
                self.refine(adjust=adjust, repeat=None)
            return

        # actual refine
        new_vertices = []
        new_vertex_indices = {}
        n_vertices = len(self.vertices)
        faces = []
        refine = np.ones(len(self.faces), dtype=bool)

        for idx, r in enumerate(refine):
            if r:
                i, j, k = self.faces[idx]
                for a, b in [(i, j), (j, k), (k, i)]:
                    ab = tuple(sorted([a, b]))
                    if ab not in new_vertex_indices:
                        new_vertices.append((self.vertices[a] + self.vertices[b]) / 2)
                        new_vertex_indices[ab] = n_vertices
                        n_vertices += 1
                for a, b, c in [
                    [i, (i, j), (i, k)],
                    [j, (j, k), (i, j)],
                    [k, (i, k), (j, k)],
                    [(i, j), (j, k), (k, i)],
                ]:
                    if isinstance(a, tuple):
                        a = new_vertex_indices[tuple(sorted(a))]
                    if isinstance(b, tuple):
                        b = new_vertex_indices[tuple(sorted(b))]
                    if isinstance(c, tuple):
                        c = new_vertex_indices[tuple(sorted(c))]
                    faces.append([a, b, c])
            else:
                faces.append(self.faces[idx])

        # adjust new vertex positions
        if adjust is None:
            try:
                adjust = self.adjust
            except AttributeError:
                pass
        if adjust is not None:
            new_vertices = [adjust(v) for v in new_vertices]

        self.faces = np.array(faces)
        self.vertices = np.concatenate([self.vertices, new_vertices])

    def plot(self, surface_kwargs=(), line_kwargs=None, vertex_kwargs=None, layout_kwargs=None, fig=None):
        data = []
        x, y, z = np.array(self.vertices).T
        i, j, k = np.array(self.faces).T

        # surface
        if surface_kwargs is not None:
            data.append(go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, **dict(surface_kwargs)))

        # lines
        if line_kwargs is not None:
            edge_x, edge_y, edge_z = [], [], []
            for from_idx, to_idx in [(i, j), (j, k), (k, i)]:
                for f, t in zip(from_idx, to_idx):
                    if f < t:
                        edge_x += [x[f], x[t], None]
                        edge_y += [y[f], y[t], None]
                        edge_z += [z[f], z[t], None]
            data.append(go.Scatter3d(x=edge_x, y=edge_y, z=edge_z, **{**dict(mode="lines"), **dict(line_kwargs)}))

        # vertices
        if vertex_kwargs is not None:
            data.append(go.Scatter3d(x=x, y=y, z=z, **{**dict(mode="markers"), **dict(vertex_kwargs)}))

        # figure
        if fig is None:
            fig = go.Figure(data=data)
        else:
            for d in data:
                fig.add_trace(d)

        # layout
        if layout_kwargs is not None:
            fig.update_layout(**layout_kwargs)
        return fig

/home/runner/work/snippets/snippets/examples/plot_3D_objects.py:34: SyntaxWarning:

invalid escape sequence '\ '

Helper function to plot surfaces

def plot_surface(surface, smooth=False, surface_kwargs=(), line_kwargs=None, vertex_kwargs=None, layout_kwargs=(), fig=None):
    _surface_kwargs = dict(color='lightblue',
                           opacity=1,
                           lightposition=dict(x=100, y=100, z=100),
                           showscale=False,
                           # 1) --- material / surface properties
                           lighting=dict(
                               ambient=0.25,  # soft overall illumination
                               diffuse=0.7,  # matte component (0-1)
                               specular=0.05,  # mirror-like shine (0-2 is safe)
                               roughness=0.4,  # width + contrast of the specular spot (0-1)
                               fresnel=0.,  # view-angle-dependent reflectance (0-5)
                               vertexnormalsepsilon=0, facenormalsepsilon=0  # avoid rendering artifacts
                           ),
                           # 2) --- position of the lamp in data coordinates
                           # lighting=dict(
                           #     ambient=0.5, diffuse=0.8, specular=0.2, roughness=0.3, fresnel=0.2,
                           #     vertexnormalsepsilon=0, facenormalsepsilon=0))
                           )
    if smooth:
        _surface_kwargs['flatshading'] = False
    else:
        _surface_kwargs['flatshading'] = True
    _line_kwargs = dict(line=dict(width=2, color="black"))
    _vertex_kwargs = dict(marker=dict(size=4, color='red'))
    _layout_kwargs = dict(scene=dict(aspectmode='data',
                                     xaxis=dict(visible=False),
                                     yaxis=dict(visible=False),
                                     zaxis=dict(visible=False),
                                     ),
                          margin=dict(l=0, r=0, b=0, t=30))

    return surface.plot(surface_kwargs={**_surface_kwargs, **dict(surface_kwargs)} if surface_kwargs is not None else None,
                        line_kwargs={**_line_kwargs, **dict(line_kwargs)} if line_kwargs is not None else None,
                        vertex_kwargs={**_vertex_kwargs, **dict(vertex_kwargs)} if vertex_kwargs is not None else None,
                        layout_kwargs={**_layout_kwargs, **dict(layout_kwargs)} if layout_kwargs is not None else None,
                        fig=fig)

Sphere

class Sphere(Surface):
    """Unit sphere based on a refined icosahedron"""

    phi = (1 + np.sqrt(5)) / 2

    def __init__(self, refine=3):
        phi = self.phi
        vertices = np.array([[-1, phi, 0],
                             [ 1,  phi,  0],
                             [-1, -phi,  0],
                             [ 1, -phi,  0],
                             [ 0, -1,  phi],
                             [ 0,  1,  phi],
                             [ 0, -1, -phi],
                             [ 0,  1, -phi],
                             [ phi,  0, -1],
                             [ phi,  0,  1],
                             [-phi,  0, -1],
                             [-phi,  0,  1]]) / np.linalg.norm([-1,  phi,  0])
        faces = [[0, 11, 5], [0, 5, 1], [0, 1, 7], [0, 7, 10], [0, 10, 11],
                 [1, 5, 9], [5, 11, 4], [11, 10, 2], [10, 7, 6], [7, 1, 8],
                 [3, 9, 4], [3, 4, 2], [3, 2, 6], [3, 6, 8], [3, 8, 9],
                 [4, 9, 5], [2, 4, 11], [6, 2, 10], [8, 6, 7], [9, 8, 1]]
        super().__init__(vertices=vertices, faces=faces)
        self.refine(repeat=refine)

    def adjust(self, v):
        v = np.array(v)
        return v / np.linalg.norm(v)

sphere = Sphere(refine=0)
plot_surface(sphere, smooth=False)

sphere = Sphere(refine=3)
plot_surface(sphere, smooth=True)

Cylinder

class Cylinder(Surface):

    def __init__(self, levels=2, bottom=True, top=True, loop=False, sangle=1, refine=3):
        assert 0 < sangle <= 1, sangle
        closed = sangle == 1

        vertices = []
        for z in np.linspace(0, 1, levels + 1):
            vertices += [[np.cos(phi) / 2, np.sin(phi) / 2, z]
                         for phi in np.linspace(start=0, stop=2 * np.pi * sangle,
                                                num=4 if closed else 5,
                                                endpoint=not closed)]

        faces = []
        for l in range(levels):
            if closed:
                (a, b, c, d,
                 e, f, g, h) = np.arange(8) + 4 * l
                a_ = a
                e_ = e
            else:
                (a, b, c, d, a_,
                 e, f, g, h, e_) = np.arange(10) + 5 * l
            faces += [[a, b, e], [b, c, f], [c, d, g], [d, a_, h]]
            faces += [[f, e, b], [g, f, c], [h, g, d], [e_, h, a_]]

        super().__init__(vertices=vertices, faces=faces)
        self.refine(repeat=refine)

        vertices = list(self.vertices)
        faces = list(self.faces)

        def angle(p, sangle=sangle):
            x, y = p[1][:2]
            angle = np.arctan2(y, x)
            return angle

        def rim(vertices, z, rev=False):
            rim = [(idx, v) for idx, v in enumerate(vertices) if v[2] == z]
            rim = sorted(rim, key=angle)
            if rev:
                rim = reversed(rim)
            return [idx for idx, v in rim]

        top_rim = rim(vertices, z=1)
        bottom_rim = rim(vertices, z=0)
        num = len(bottom_rim)
        assert num == len(top_rim)
        indices_from = np.arange(num)
        indices_to = (np.arange(num) + 1) % num  # loop around rim to close
        if not closed:
            # exclude index pairs across the opening (before refining)
            top_indices = [(i, j) for i, j in zip(indices_from, indices_to) if (top_rim[i], top_rim[j]) != (4 + 5 * levels, 5 * levels)]
            bottom_indices = [(i, j) for i, j in zip(indices_from, indices_to) if (bottom_rim[i], bottom_rim[j]) != (4, 0)]
        else:
            top_indices = zip(indices_from, indices_to)
            bottom_indices = zip(indices_from, indices_to)
        if top:
            centre_idx = len(vertices)
            vertices.append([0., 0., 1.])
            faces += [[top_rim[i], top_rim[j], centre_idx] for i, j in top_indices]
        if bottom:
            centre_idx = len(vertices)
            vertices.append([0., 0., 0.])
            faces += [[bottom_rim[j], bottom_rim[i], centre_idx] for i, j in bottom_indices]
        if loop:
            faces += [[top_rim[i], top_rim[j], bottom_rim[i]] for i, j in top_indices]
            faces += [[bottom_rim[j], bottom_rim[i], top_rim[j]] for i, j in bottom_indices]

        self.vertices = np.array(vertices)
        self.faces = np.array(faces)

    def adjust(self, v):
        x, y, z = v
        xy = np.array([x, y])
        r = np.linalg.norm(xy)
        if r == 0:
            return v
        else:
            return 0.5 * x / r, 0.5 * y / r, z


plot_surface(Cylinder(), line_kwargs=(), vertex_kwargs=())

Torus

class Torus(Surface):

    @classmethod
    def torus(cls, phi1, phi2, r1, r2):
        """

        :param phi1: first angle in [0, 1] (for "large" circle)
        :param phi2: second angle in [0, 1] (for "small" circle)
        :param r1: first radius (for "larger" circle)
        :param r2: second radius (for "small" circle)
        :return:
        """
        x = np.cos(2 * np.pi * phi1) * (r1 + np.cos(2 * np.pi * phi2) * r2)
        y = np.sin(2 * np.pi * phi1) * (r1 + np.cos(2 * np.pi * phi2) * r2)
        z = np.sin(2 * np.pi * phi2) * r2
        return x, y, z

    def __init__(self, R, r, levels=5, refine=4):
        s = Cylinder(levels=levels, bottom=False, top=False, loop=True, refine=refine)
        x, y, z = s.vertices.T
        levels = levels * 2 ** refine
        z *= levels / (levels + 1)
        theta = np.arctan2(y, x) / (2 * np.pi)
        x, y, z = self.torus(phi1=z, phi2=theta, r1=R, r2=r)
        super().__init__(vertices=np.stack([x, y, z], axis=-1), faces=s.faces)

torus = Torus(R=1, r=0.3)
plot_surface(torus, line_kwargs=(), smooth=False)

plot_surface(torus, smooth=True)

Torus

class Torus(Surface):

    @classmethod
    def torus(cls, phi1, phi2, r1, r2):
        """

        :param phi1: first angle in [0, 1] (for "large" circle)
        :param phi2: second angle in [0, 1] (for "small" circle)
        :param r1: first radius (for "larger" circle)
        :param r2: second radius (for "small" circle)
        :return:
        """
        x = np.cos(2 * np.pi * phi1) * (r1 + np.cos(2 * np.pi * phi2) * r2)
        y = np.sin(2 * np.pi * phi1) * (r1 + np.cos(2 * np.pi * phi2) * r2)
        z = np.sin(2 * np.pi * phi2) * r2
        return x, y, z

    def __init__(self, R, r, levels=5, refine=4, sphi1=1, sphi2=1, dphi1=0, dphi2=0, bottom=False, top=False, loop=None):
        assert 0 < sphi1 <= 1, sphi1
        assert 0 < sphi2 <= 1, sphi2
        if loop is None:
            loop = sphi1 == 1
        s = Cylinder(levels=levels, bottom=bottom, top=top, loop=loop, sangle=sphi2, refine=refine)
        x, y, z = s.vertices.T
        levels = levels * 2 ** refine
        z *= levels / (levels + 1)
        theta = np.arctan2(y, x) / (2 * np.pi)
        r = np.sqrt(x ** 2 + y ** 2) * 2 * r
        x, y, z = self.torus(phi1=z * sphi1 + dphi1, phi2=theta + dphi2, r1=R, r2=r)
        super().__init__(vertices=np.stack([x, y, z], axis=-1), faces=s.faces)

torus = Torus(R=1, r=0.3)
plot_surface(torus, line_kwargs=(), smooth=False)

show(plot_surface(torus, smooth=True), __name__)