vedo.transformations

   1#!/usr/bin/env python3
   2# -*- coding: utf-8 -*-
   3from typing import List
   4from typing_extensions import Self
   5import numpy as np
   6
   7import vedo.vtkclasses as vtki # a wrapper for lazy imports
   8
   9__docformat__ = "google"
  10
  11__doc__ = """
  12Submodule to work with linear and non-linear transformations<br>
  13
  14![](https://vedo.embl.es/images/feats/transforms.png)
  15"""
  16
  17__all__ = [
  18    "LinearTransform",
  19    "NonLinearTransform",
  20    "TransformInterpolator",
  21    "spher2cart",
  22    "cart2spher",
  23    "cart2cyl",
  24    "cyl2cart",
  25    "cyl2spher",
  26    "spher2cyl",
  27    "cart2pol",
  28    "pol2cart",
  29]
  30
  31###################################################
  32def _is_sequence(arg):
  33    if hasattr(arg, "strip"):
  34        return False
  35    if hasattr(arg, "__getslice__"):
  36        return True
  37    if hasattr(arg, "__iter__"):
  38        return True
  39    return False
  40
  41
  42###################################################
  43class LinearTransform:
  44    """Work with linear transformations."""
  45
  46    def __init__(self, T=None) -> None:
  47        """
  48        Define a linear transformation.
  49        Can be saved to file and reloaded.
  50
  51        Arguments:
  52            T : (str, vtkTransform, numpy array)
  53                input transformation. Defaults to unit.
  54
  55        Example:
  56            ```python
  57            from vedo import *
  58            settings.use_parallel_projection = True
  59
  60            LT = LinearTransform()
  61            LT.translate([3,0,1]).rotate_z(45)
  62            LT.comment = "shifting by (3,0,1) and rotating by 45 deg"
  63            print(LT)
  64
  65            sph = Sphere(r=0.2)
  66            sph.apply_transform(LT) # same as: LT.move(s1)
  67            print(sph.transform)
  68
  69            show(Point([0,0,0]), sph, str(LT.matrix), axes=1).close()
  70            ```
  71        """
  72        self.name = "LinearTransform"
  73        self.filename = ""
  74        self.comment = ""
  75
  76        if T is None:
  77            T = vtki.vtkTransform()
  78
  79        elif isinstance(T, vtki.vtkMatrix4x4):
  80            S = vtki.vtkTransform()
  81            S.SetMatrix(T)
  82            T = S
  83
  84        elif isinstance(T, vtki.vtkLandmarkTransform):
  85            S = vtki.vtkTransform()
  86            S.SetMatrix(T.GetMatrix())
  87            T = S
  88
  89        elif _is_sequence(T):
  90            S = vtki.vtkTransform()
  91            M = vtki.vtkMatrix4x4()
  92            n = len(T)
  93            for i in range(n):
  94                for j in range(n):
  95                    M.SetElement(i, j, T[i][j])
  96            S.SetMatrix(M)
  97            T = S
  98
  99        elif isinstance(T, vtki.vtkLinearTransform):
 100            S = vtki.vtkTransform()
 101            S.DeepCopy(T)
 102            T = S
 103
 104        elif isinstance(T, LinearTransform):
 105            S = vtki.vtkTransform()
 106            S.DeepCopy(T.T)
 107            T = S
 108
 109        elif isinstance(T, str):
 110            import json
 111            self.filename = str(T)
 112            try:
 113                with open(self.filename, "r") as read_file:
 114                    D = json.load(read_file)
 115                self.name = D["name"]
 116                self.comment = D["comment"]
 117                matrix = np.array(D["matrix"])
 118            except json.decoder.JSONDecodeError:
 119                ### assuming legacy vedo format E.g.:
 120                # aligned by manual_align.py
 121                # 0.8026854838223 -0.0789823873914 -0.508476844097  38.17377632072
 122                # 0.0679734082661  0.9501827489452 -0.040289803376 -69.53864247951
 123                # 0.5100652300642 -0.0023313569781  0.805555043665 -81.20317788519
 124                # 0.0 0.0 0.0 1.0
 125                with open(self.filename, "r", encoding="UTF-8") as read_file:
 126                    lines = read_file.readlines()
 127                    i = 0
 128                    matrix = np.eye(4)
 129                    for l in lines:
 130                        if l.startswith("#"):
 131                            self.comment = l.replace("#", "").strip()
 132                            continue
 133                        vals = l.split(" ")
 134                        for j in range(len(vals)):
 135                            v = vals[j].replace("\n", "")
 136                            if v != "":
 137                                matrix[i, j] = float(v)
 138                        i += 1
 139            T = vtki.vtkTransform()
 140            m = vtki.vtkMatrix4x4()
 141            for i in range(4):
 142                for j in range(4):
 143                    m.SetElement(i, j, matrix[i][j])
 144            T.SetMatrix(m)
 145
 146        self.T = T
 147        self.T.PostMultiply()
 148        self.inverse_flag = False
 149
 150    def __str__(self):
 151        module = self.__class__.__module__
 152        name = self.__class__.__name__
 153        s = f"\x1b[7m\x1b[1m{module}.{name} at ({hex(id(self))})".ljust(75) + "\x1b[0m"
 154        s += "\nname".ljust(15) + ": " + self.name
 155        if self.filename:
 156            s += "\nfilename".ljust(15) + ": " + self.filename
 157        if self.comment:
 158            s += "\ncomment".ljust(15) + f': \x1b[3m"{self.comment}"\x1b[0m'
 159        s += f"\nconcatenations".ljust(15) + f": {self.ntransforms}"
 160        s += "\ninverse flag".ljust(15) + f": {bool(self.inverse_flag)}"
 161        arr = np.array2string(self.matrix,
 162            separator=', ', precision=6, suppress_small=True)
 163        s += "\nmatrix 4x4".ljust(15) + f":\n{arr}"
 164        return s
 165
 166    def __repr__(self):
 167        return self.__str__()
 168
 169    def print(self) -> "LinearTransform":
 170        """Print transformation."""
 171        print(self.__str__())
 172        return self
 173
 174    def __call__(self, obj):
 175        """
 176        Apply transformation to object or single point.
 177        Same as `move()` except that a copy is returned.
 178        """
 179        return self.move(obj.copy())
 180    
 181    def transform_point(self, p) -> np.ndarray:
 182        """
 183        Apply transformation to a single point.
 184        """
 185        if len(p) == 2:
 186            p = [p[0], p[1], 0]
 187        return np.array(self.T.TransformFloatPoint(p))
 188
 189    def move(self, obj):
 190        """
 191        Apply transformation to object or single point.
 192
 193        Note:
 194            When applying a transformation to a mesh, the mesh is modified in place.
 195            If you want to keep the original mesh unchanged, use `clone()` method.
 196
 197        Example:
 198            ```python
 199            from vedo import *
 200            settings.use_parallel_projection = True
 201
 202            LT = LinearTransform()
 203            LT.translate([3,0,1]).rotate_z(45)
 204            print(LT)
 205
 206            s = Sphere(r=0.2)
 207            LT.move(s)
 208            # same as:
 209            # s.apply_transform(LT)
 210
 211            zero = Point([0,0,0])
 212            show(s, zero, axes=1).close()
 213            ```
 214        """
 215        if _is_sequence(obj):
 216            n = len(obj)
 217            if n == 2:
 218                obj = [obj[0], obj[1], 0]
 219            return np.array(self.T.TransformFloatPoint(obj))
 220
 221        obj.apply_transform(self)
 222        return obj
 223
 224    def reset(self) -> Self:
 225        """Reset transformation."""
 226        self.T.Identity()
 227        return self
 228    
 229    def compute_main_axes(self) -> np.ndarray:
 230        """
 231        Compute main axes of the transformation matrix.
 232        These are the axes of the ellipsoid that is the 
 233        image of the unit sphere under the transformation.
 234
 235        Example:
 236        ```python
 237        from vedo import *
 238        settings.use_parallel_projection = True
 239
 240        M = np.random.rand(3,3)-0.5
 241        print(M)
 242        print(" M@[1,0,0] =", M@[1,1,0])
 243
 244        ######################
 245        A = LinearTransform(M)
 246        print(A)
 247        pt = Point([1,1,0])
 248        print(A(pt).vertices[0], "is the same as", A([1,1,0]))
 249
 250        maxes = A.compute_main_axes()
 251
 252        arr1 = Arrow([0,0,0], maxes[0]).c('r')
 253        arr2 = Arrow([0,0,0], maxes[1]).c('g')
 254        arr3 = Arrow([0,0,0], maxes[2]).c('b')
 255
 256        sphere1 = Sphere().wireframe().lighting('off')
 257        sphere1.cmap('hot', sphere1.vertices[:,2])
 258
 259        sphere2 = sphere1.clone().apply_transform(A)
 260
 261        show([sphere1, [sphere2, arr1, arr2, arr3]], N=2, axes=1, bg='bb')
 262        ```
 263        """
 264        m = self.matrix3x3
 265        eigval, eigvec = np.linalg.eig(m @ m.T)
 266        eigval = np.sqrt(eigval)
 267        return  np.array([
 268            eigvec[:,0] * eigval[0],
 269            eigvec[:,1] * eigval[1],
 270            eigvec[:,2] * eigval[2],
 271        ])
 272
 273    def pop(self) -> Self:
 274        """Delete the transformation on the top of the stack
 275        and sets the top to the next transformation on the stack."""
 276        self.T.Pop()
 277        return self
 278
 279    def is_identity(self) -> bool:
 280        """Check if the transformation is the identity."""
 281        m = self.T.GetMatrix()
 282        M = [[m.GetElement(i, j) for j in range(4)] for i in range(4)]
 283        if np.allclose(M - np.eye(4), 0):
 284            return True
 285        return False
 286
 287    def invert(self) -> Self:
 288        """Invert the transformation. Acts in-place."""
 289        self.T.Inverse()
 290        self.inverse_flag = bool(self.T.GetInverseFlag())
 291        return self
 292
 293    def compute_inverse(self) -> "LinearTransform":
 294        """Compute the inverse."""
 295        t = self.clone()
 296        t.invert()
 297        return t
 298
 299    def transpose(self) -> Self:
 300        """Transpose the transformation. Acts in-place."""
 301        M = vtki.vtkMatrix4x4()
 302        self.T.GetTranspose(M)
 303        self.T.SetMatrix(M)
 304        return self
 305
 306    def copy(self) -> "LinearTransform":
 307        """Return a copy of the transformation. Alias of `clone()`."""
 308        return self.clone()
 309
 310    def clone(self) -> "LinearTransform":
 311        """Clone transformation to make an exact copy."""
 312        return LinearTransform(self.T)
 313
 314    def concatenate(self, T, pre_multiply=False) -> Self:
 315        """
 316        Post-multiply (by default) 2 transfomations.
 317        T can also be a 4x4 matrix or 3x3 matrix.
 318
 319        Example:
 320            ```python
 321            from vedo import LinearTransform
 322
 323            A = LinearTransform()
 324            A.rotate_x(45)
 325            A.translate([7,8,9])
 326            A.translate([10,10,10])
 327            A.name = "My transformation A"
 328            print(A)
 329
 330            B = A.compute_inverse()
 331            B.shift([1,2,3])
 332            B.name = "My transformation B (shifted inverse of A)"
 333            print(B)
 334
 335            # A is applied first, then B
 336            # print("A.concatenate(B)", A.concatenate(B))
 337
 338            # B is applied first, then A
 339            print(B*A)
 340            ```
 341        """
 342        if _is_sequence(T):
 343            S = vtki.vtkTransform()
 344            M = vtki.vtkMatrix4x4()
 345            n = len(T)
 346            for i in range(n):
 347                for j in range(n):
 348                    M.SetElement(i, j, T[i][j])
 349            S.SetMatrix(M)
 350            T = S
 351
 352        if pre_multiply:
 353            self.T.PreMultiply()
 354        try:
 355            self.T.Concatenate(T)
 356        except:
 357            self.T.Concatenate(T.T)
 358        self.T.PostMultiply()
 359        return self
 360
 361    def __mul__(self, A):
 362        """Pre-multiply 2 transfomations."""
 363        return self.concatenate(A, pre_multiply=True)
 364
 365    def get_concatenated_transform(self, i) -> "LinearTransform":
 366        """Get intermediate matrix by concatenation index."""
 367        return LinearTransform(self.T.GetConcatenatedTransform(i))
 368
 369    @property
 370    def ntransforms(self) -> int:
 371        """Get the number of concatenated transforms."""
 372        return self.T.GetNumberOfConcatenatedTransforms()
 373
 374    def translate(self, p) -> Self:
 375        """Translate, same as `shift`."""
 376        if len(p) == 2:
 377            p = [p[0], p[1], 0]
 378        self.T.Translate(p)
 379        return self
 380
 381    def shift(self, p) -> Self:
 382        """Shift, same as `translate`."""
 383        return self.translate(p)
 384
 385    def scale(self, s, origin=True) -> Self:
 386        """Scale."""
 387        if not _is_sequence(s):
 388            s = [s, s, s]
 389
 390        if origin is True:
 391            p = np.array(self.T.GetPosition())
 392            if np.linalg.norm(p) > 0:
 393                self.T.Translate(-p)
 394                self.T.Scale(*s)
 395                self.T.Translate(p)
 396            else:
 397                self.T.Scale(*s)
 398
 399        elif _is_sequence(origin):
 400            origin = np.asarray(origin)
 401            self.T.Translate(-origin)
 402            self.T.Scale(*s)
 403            self.T.Translate(origin)
 404
 405        else:
 406            self.T.Scale(*s)
 407        return self
 408
 409    def rotate(self, angle, axis=(1, 0, 0), point=(0, 0, 0), rad=False) -> Self:
 410        """
 411        Rotate around an arbitrary `axis` passing through `point`.
 412
 413        Example:
 414            ```python
 415            from vedo import *
 416            c1 = Cube()
 417            c2 = c1.clone().c('violet').alpha(0.5) # copy of c1
 418            v = vector(0.2, 1, 0)
 419            p = vector(1.0, 0, 0)  # axis passes through this point
 420            c2.rotate(90, axis=v, point=p)
 421            l = Line(p-v, p+v).c('red5').lw(3)
 422            show(c1, l, c2, axes=1).close()
 423            ```
 424            ![](https://vedo.embl.es/images/feats/rotate_axis.png)
 425        """
 426        if np.all(axis == 0):
 427            return self
 428        if not angle:
 429            return self
 430        if rad:
 431            anglerad = angle
 432        else:
 433            anglerad = np.deg2rad(angle)
 434        
 435        axis = np.asarray(axis) / np.linalg.norm(axis)
 436        a = np.cos(anglerad / 2)
 437        b, c, d = -axis * np.sin(anglerad / 2)
 438        aa, bb, cc, dd = a * a, b * b, c * c, d * d
 439        bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
 440        R = np.array(
 441            [
 442                [aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
 443                [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
 444                [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc],
 445            ]
 446        )
 447        rv = np.dot(R, self.T.GetPosition() - np.asarray(point)) + point
 448
 449        if rad:
 450            angle *= 180.0 / np.pi
 451        # this vtk method only rotates in the origin of the object:
 452        self.T.RotateWXYZ(angle, axis[0], axis[1], axis[2])
 453        self.T.Translate(rv - np.array(self.T.GetPosition()))
 454        return self
 455
 456    def _rotatexyz(self, axe, angle, rad, around):
 457        if not angle:
 458            return self
 459        if rad:
 460            angle *= 180 / np.pi
 461
 462        rot = dict(x=self.T.RotateX, y=self.T.RotateY, z=self.T.RotateZ)
 463
 464        if around is None:
 465            # rotate around its origin
 466            rot[axe](angle)
 467        else:
 468            # displacement needed to bring it back to the origin
 469            self.T.Translate(-np.asarray(around))
 470            rot[axe](angle)
 471            self.T.Translate(around)
 472        return self
 473
 474    def rotate_x(self, angle: float, rad=False, around=None) -> Self:
 475        """
 476        Rotate around x-axis. If angle is in radians set `rad=True`.
 477
 478        Use `around` to define a pivoting point.
 479        """
 480        return self._rotatexyz("x", angle, rad, around)
 481
 482    def rotate_y(self, angle: float, rad=False, around=None) -> Self:
 483        """
 484        Rotate around y-axis. If angle is in radians set `rad=True`.
 485
 486        Use `around` to define a pivoting point.
 487        """
 488        return self._rotatexyz("y", angle, rad, around)
 489
 490    def rotate_z(self, angle: float, rad=False, around=None) -> Self:
 491        """
 492        Rotate around z-axis. If angle is in radians set `rad=True`.
 493
 494        Use `around` to define a pivoting point.
 495        """
 496        return self._rotatexyz("z", angle, rad, around)
 497
 498    def set_position(self, p) -> Self:
 499        """Set position."""
 500        if len(p) == 2:
 501            p = np.array([p[0], p[1], 0])
 502        q = np.array(self.T.GetPosition())
 503        self.T.Translate(p - q)
 504        return self
 505
 506    # def set_scale(self, s):
 507    #     """Set absolute scale."""
 508    #     if not _is_sequence(s):
 509    #         s = [s, s, s]
 510    #     s0, s1, s2 = 1, 1, 1
 511    #     b = self.T.GetScale()
 512    #     print(b)
 513    #     if b[0]:
 514    #         s0 = s[0] / b[0]
 515    #     if b[1]:
 516    #         s1 = s[1] / b[1]
 517    #     if b[2]:
 518    #         s2 = s[2] / b[2]
 519    #     self.T.Scale(s0, s1, s2)
 520    #     print()
 521    #     return self
 522
 523    def get_scale(self) -> np.ndarray:
 524        """Get current scale."""
 525        return np.array(self.T.GetScale())
 526
 527    @property
 528    def orientation(self) -> np.ndarray:
 529        """Compute orientation."""
 530        return np.array(self.T.GetOrientation())
 531
 532    @property
 533    def position(self) -> np.ndarray:
 534        """Compute position."""
 535        return np.array(self.T.GetPosition())
 536
 537    @property
 538    def matrix(self) -> np.ndarray:
 539        """Get the 4x4 trasformation matrix."""
 540        m = self.T.GetMatrix()
 541        M = [[m.GetElement(i, j) for j in range(4)] for i in range(4)]
 542        return np.array(M)
 543
 544    @matrix.setter
 545    def matrix(self, M) -> None:
 546        """Set trasformation by assigning a 4x4 or 3x3 numpy matrix."""
 547        n = len(M)
 548        m = vtki.vtkMatrix4x4()
 549        for i in range(n):
 550            for j in range(n):
 551                m.SetElement(i, j, M[i][j])
 552        self.T.SetMatrix(m)
 553
 554    @property
 555    def matrix3x3(self) -> np.ndarray:
 556        """Get the 3x3 trasformation matrix."""
 557        m = self.T.GetMatrix()
 558        M = [[m.GetElement(i, j) for j in range(3)] for i in range(3)]
 559        return np.array(M)
 560
 561    def write(self, filename="transform.mat") -> Self:
 562        """Save transformation to ASCII file."""
 563        import json
 564        m = self.T.GetMatrix()
 565        M = [[m.GetElement(i, j) for j in range(4)] for i in range(4)]
 566        arr = np.array(M)
 567        dictionary = {
 568            "name": self.name,
 569            "comment": self.comment,
 570            "matrix": arr.astype(float).tolist(),
 571            "ntransforms": self.ntransforms,
 572        }
 573        with open(filename, "w") as outfile:
 574            json.dump(dictionary, outfile, sort_keys=True, indent=2)
 575        return self
 576
 577    def reorient(
 578        self, initaxis, newaxis, around=(0, 0, 0), rotation=0.0, rad=False, xyplane=True
 579    ) -> Self:
 580        """
 581        Set/Get object orientation.
 582
 583        Arguments:
 584            rotation : (float)
 585                rotate object around newaxis.
 586            concatenate : (bool)
 587                concatenate the orientation operation with the previous existing transform (if any)
 588            rad : (bool)
 589                set to True if angle is expressed in radians.
 590            xyplane : (bool)
 591                make an extra rotation to keep the object aligned to the xy-plane
 592        """
 593        newaxis = np.asarray(newaxis) / np.linalg.norm(newaxis)
 594        initaxis = np.asarray(initaxis) / np.linalg.norm(initaxis)
 595
 596        if not np.any(initaxis - newaxis):
 597            return self
 598
 599        if not np.any(initaxis + newaxis):
 600            print("Warning: in reorient() initaxis and newaxis are parallel")
 601            newaxis += np.array([0.0000001, 0.0000002, 0.0])
 602            angleth = np.pi
 603        else:
 604            angleth = np.arccos(np.dot(initaxis, newaxis))
 605        crossvec = np.cross(initaxis, newaxis)
 606
 607        p = np.asarray(around)
 608        self.T.Translate(-p)
 609        if rotation:
 610            if rad:
 611                rotation = np.rad2deg(rotation)
 612            self.T.RotateWXYZ(rotation, initaxis)
 613
 614        self.T.RotateWXYZ(np.rad2deg(angleth), crossvec)
 615
 616        if xyplane:
 617            self.T.RotateWXYZ(-self.orientation[0] * 1.4142, newaxis)
 618
 619        self.T.Translate(p)
 620        return self
 621
 622
 623###################################################
 624class NonLinearTransform:
 625    """Work with non-linear transformations."""
 626
 627    def __init__(self, T=None, **kwargs) -> None:
 628        """
 629        Define a non-linear transformation.
 630        Can be saved to file and reloaded.
 631
 632        Arguments:
 633            T : (vtkThinPlateSplineTransform, str, dict)
 634                vtk transformation.
 635                If T is a string, it is assumed to be a filename.
 636                If T is a dictionary, it is assumed to be a set of keyword arguments.
 637                Defaults to None.
 638            **kwargs : (dict)
 639                keyword arguments to define the transformation.
 640                The following keywords are accepted:
 641                - name : (str) name of the transformation
 642                - comment : (str) comment
 643                - source_points : (list) source points
 644                - target_points : (list) target points
 645                - mode : (str) either '2d' or '3d'
 646                - sigma : (float) sigma parameter
 647
 648        Example:
 649            ```python
 650            from vedo import *
 651            settings.use_parallel_projection = True
 652
 653            NLT = NonLinearTransform()
 654            NLT.source_points = [[-2,0,0], [1,2,1], [2,-2,2]]
 655            NLT.target_points = NLT.source_points + np.random.randn(3,3)*0.5
 656            NLT.mode = '3d'
 657            print(NLT)
 658
 659            s1 = Sphere()
 660            NLT.move(s1)
 661            # same as:
 662            # s1.apply_transform(NLT)
 663
 664            arrs = Arrows(NLT.source_points, NLT.target_points)
 665            show(s1, arrs, Sphere().alpha(0.1), axes=1).close()
 666            ```
 667        """
 668
 669        self.name = "NonLinearTransform"
 670        self.filename = ""
 671        self.comment = ""
 672
 673        if T is None and len(kwargs) == 0:
 674            T = vtki.vtkThinPlateSplineTransform()
 675
 676        elif isinstance(T, vtki.vtkThinPlateSplineTransform):
 677            S = vtki.vtkThinPlateSplineTransform()
 678            S.DeepCopy(T)
 679            T = S
 680
 681        elif isinstance(T, NonLinearTransform):
 682            S = vtki.vtkThinPlateSplineTransform()
 683            S.DeepCopy(T.T)
 684            T = S
 685
 686        elif isinstance(T, str):
 687            import json
 688            filename = str(T)
 689            self.filename = filename
 690            with open(filename, "r") as read_file:
 691                D = json.load(read_file)
 692            self.name = D["name"]
 693            self.comment = D["comment"]
 694            source = D["source_points"]
 695            target = D["target_points"]
 696            mode = D["mode"]
 697            sigma = D["sigma"]
 698
 699            T = vtki.vtkThinPlateSplineTransform()
 700            vptss = vtki.vtkPoints()
 701            for p in source:
 702                if len(p) == 2:
 703                    p = [p[0], p[1], 0.0]
 704                vptss.InsertNextPoint(p)
 705            T.SetSourceLandmarks(vptss)
 706            vptst = vtki.vtkPoints()
 707            for p in target:
 708                if len(p) == 2:
 709                    p = [p[0], p[1], 0.0]
 710                vptst.InsertNextPoint(p)
 711            T.SetTargetLandmarks(vptst)
 712            T.SetSigma(sigma)
 713            if mode == "2d":
 714                T.SetBasisToR2LogR()
 715            elif mode == "3d":
 716                T.SetBasisToR()
 717            else:
 718                print(f'In {filename} mode can be either "2d" or "3d"')
 719
 720        elif len(kwargs) > 0:
 721            T = kwargs.copy()
 722            self.name = T.pop("name", "NonLinearTransform")
 723            self.comment = T.pop("comment", "")
 724            source = T.pop("source_points", [])
 725            target = T.pop("target_points", [])
 726            mode = T.pop("mode", "3d")
 727            sigma = T.pop("sigma", 1.0)
 728            if len(T) > 0:
 729                print("Warning: NonLinearTransform got unexpected keyword arguments:")
 730                print(T)
 731
 732            T = vtki.vtkThinPlateSplineTransform()
 733            vptss = vtki.vtkPoints()
 734            for p in source:
 735                if len(p) == 2:
 736                    p = [p[0], p[1], 0.0]
 737                vptss.InsertNextPoint(p)
 738            T.SetSourceLandmarks(vptss)
 739            vptst = vtki.vtkPoints()
 740            for p in target:
 741                if len(p) == 2:
 742                    p = [p[0], p[1], 0.0]
 743                vptst.InsertNextPoint(p)
 744            T.SetTargetLandmarks(vptst)
 745            T.SetSigma(sigma)
 746            if mode == "2d":
 747                T.SetBasisToR2LogR()
 748            elif mode == "3d":
 749                T.SetBasisToR()
 750            else:
 751                print(f'Warning: mode can be either "2d" or "3d"')
 752
 753        self.T = T
 754        self.inverse_flag = False
 755
 756    def __str__(self):
 757        module = self.__class__.__module__
 758        name = self.__class__.__name__
 759        s = f"\x1b[7m\x1b[1m{module}.{name} at ({hex(id(self))})".ljust(75) + "\x1b[0m\n"
 760        s += "name".ljust(9) + ": "  + self.name + "\n"
 761        if self.filename:
 762            s += "filename".ljust(9) + ": " + self.filename + "\n"
 763        if self.comment:
 764            s += "comment".ljust(9) + f': \x1b[3m"{self.comment}"\x1b[0m\n'
 765        s += f"mode".ljust(9)  + f": {self.mode}\n"
 766        s += f"sigma".ljust(9) + f": {self.sigma}\n"
 767        p = self.source_points
 768        q = self.target_points
 769        s += f"sources".ljust(9) + f": {p.size}, bounds {np.min(p, axis=0)}, {np.max(p, axis=0)}\n"
 770        s += f"targets".ljust(9) + f": {q.size}, bounds {np.min(q, axis=0)}, {np.max(q, axis=0)}"
 771        return s
 772
 773    def __repr__(self):
 774        return self.__str__()
 775
 776    def print(self) -> Self:
 777        """Print transformation."""
 778        print(self.__str__())
 779        return self
 780
 781    def update(self) -> Self:
 782        """Update transformation."""
 783        self.T.Update()
 784        return self
 785
 786    @property
 787    def position(self) -> np.ndarray:
 788        """
 789        Trying to get the position of a `NonLinearTransform` always returns [0,0,0].
 790        """
 791        return np.array([0.0, 0.0, 0.0], dtype=np.float32)
 792
 793    # @position.setter
 794    # def position(self, p):
 795    #     """
 796    #     Trying to set position of a `NonLinearTransform`
 797    #     has no effect and prints a warning.
 798
 799    #     Use clone() method to create a copy of the object,
 800    #     or reset it with 'object.transform = vedo.LinearTransform()'
 801    #     """
 802    #     print("Warning: NonLinearTransform has no position.")
 803    #     print("  Use clone() method to create a copy of the object,")
 804    #     print("  or reset it with 'object.transform = vedo.LinearTransform()'")
 805
 806    @property
 807    def source_points(self) -> np.ndarray:
 808        """Get the source points."""
 809        pts = self.T.GetSourceLandmarks()
 810        vpts = []
 811        if pts:
 812            for i in range(pts.GetNumberOfPoints()):
 813                vpts.append(pts.GetPoint(i))
 814        return np.array(vpts, dtype=np.float32)
 815
 816    @source_points.setter
 817    def source_points(self, pts):
 818        """Set source points."""
 819        if _is_sequence(pts):
 820            pass
 821        else:
 822            pts = pts.vertices
 823        vpts = vtki.vtkPoints()
 824        for p in pts:
 825            if len(p) == 2:
 826                p = [p[0], p[1], 0.0]
 827            vpts.InsertNextPoint(p)
 828        self.T.SetSourceLandmarks(vpts)
 829
 830    @property
 831    def target_points(self) -> np.ndarray:
 832        """Get the target points."""
 833        pts = self.T.GetTargetLandmarks()
 834        vpts = []
 835        for i in range(pts.GetNumberOfPoints()):
 836            vpts.append(pts.GetPoint(i))
 837        return np.array(vpts, dtype=np.float32)
 838
 839    @target_points.setter
 840    def target_points(self, pts):
 841        """Set target points."""
 842        if _is_sequence(pts):
 843            pass
 844        else:
 845            pts = pts.vertices
 846        vpts = vtki.vtkPoints()
 847        for p in pts:
 848            if len(p) == 2:
 849                p = [p[0], p[1], 0.0]
 850            vpts.InsertNextPoint(p)
 851        self.T.SetTargetLandmarks(vpts)
 852
 853
 854    @property
 855    def sigma(self) -> float:
 856        """Set sigma."""
 857        return self.T.GetSigma()
 858
 859    @sigma.setter
 860    def sigma(self, s):
 861        """Get sigma."""
 862        self.T.SetSigma(s)
 863
 864    @property
 865    def mode(self) -> str:
 866        """Get mode."""
 867        m = self.T.GetBasis()
 868        # print("T.GetBasis()", m, self.T.GetBasisAsString())
 869        if m == 2:
 870            return "2d"
 871        elif m == 1:
 872            return "3d"
 873        else:
 874            print("Warning: NonLinearTransform has no valid mode.")
 875            return ""
 876
 877    @mode.setter
 878    def mode(self, m):
 879        """Set mode."""
 880        if m == "3d":
 881            self.T.SetBasisToR()
 882        elif m == "2d":
 883            self.T.SetBasisToR2LogR()
 884        else:
 885            print('In NonLinearTransform mode can be either "2d" or "3d"')
 886
 887    def clone(self) -> "NonLinearTransform":
 888        """Clone transformation to make an exact copy."""
 889        return NonLinearTransform(self.T)
 890
 891    def write(self, filename) -> Self:
 892        """Save transformation to ASCII file."""
 893        import json
 894
 895        dictionary = {
 896            "name": self.name,
 897            "comment": self.comment,
 898            "mode": self.mode,
 899            "sigma": self.sigma,
 900            "source_points": self.source_points.astype(float).tolist(),
 901            "target_points": self.target_points.astype(float).tolist(),
 902        }
 903        with open(filename, "w") as outfile:
 904            json.dump(dictionary, outfile, sort_keys=True, indent=2)
 905        return self
 906
 907    def invert(self) -> "NonLinearTransform":
 908        """Invert transformation."""
 909        self.T.Inverse()
 910        self.inverse_flag = bool(self.T.GetInverseFlag())
 911        return self
 912
 913    def compute_inverse(self) -> Self:
 914        """Compute inverse."""
 915        t = self.clone()
 916        t.invert()
 917        return t
 918
 919    def __call__(self, obj):
 920        """
 921        Apply transformation to object or single point.
 922        Same as `move()` except that a copy is returned.
 923        """
 924        # use copy here not clone in case user passes a numpy array
 925        return self.move(obj.copy())
 926
 927    def compute_main_axes(self, pt=(0,0,0), ds=1) -> np.ndarray:
 928        """
 929        Compute main axes of the transformation.
 930        These are the axes of the ellipsoid that is the 
 931        image of the unit sphere under the transformation.
 932
 933        Arguments:
 934            pt : (list)
 935                point to compute the axes at.
 936            ds : (float)
 937                step size to compute the axes.
 938        """
 939        if len(pt) == 2:
 940            pt = [pt[0], pt[1], 0]
 941        pt = np.asarray(pt)
 942        m = np.array([
 943            self.move(pt + [ds,0,0]),
 944            self.move(pt + [0,ds,0]),
 945            self.move(pt + [0,0,ds]),
 946        ])
 947        eigval, eigvec = np.linalg.eig(m @ m.T)
 948        eigval = np.sqrt(eigval)
 949        return np.array([
 950            eigvec[:, 0] * eigval[0],
 951            eigvec[:, 1] * eigval[1],
 952            eigvec[:, 2] * eigval[2],
 953        ])
 954
 955    def transform_point(self, p) -> np.ndarray:
 956        """
 957        Apply transformation to a single point.
 958        """
 959        if len(p) == 2:
 960            p = [p[0], p[1], 0]
 961        return np.array(self.T.TransformFloatPoint(p))
 962
 963    def move(self, obj):
 964        """
 965        Apply transformation to the argument object.
 966
 967        Note:
 968            When applying a transformation to a mesh, the mesh is modified in place.
 969            If you want to keep the original mesh unchanged, use the `clone()` method.
 970
 971        Example:
 972            ```python
 973            from vedo import *
 974            np.random.seed(0)
 975            settings.use_parallel_projection = True
 976
 977            NLT = NonLinearTransform()
 978            NLT.source_points = [[-2,0,0], [1,2,1], [2,-2,2]]
 979            NLT.target_points = NLT.source_points + np.random.randn(3,3)*0.5
 980            NLT.mode = '3d'
 981            print(NLT)
 982
 983            s1 = Sphere()
 984            NLT.move(s1)
 985            # same as:
 986            # s1.apply_transform(NLT)
 987
 988            arrs = Arrows(NLT.source_points, NLT.target_points)
 989            show(s1, arrs, Sphere().alpha(0.1), axes=1).close()
 990            ```
 991        """
 992        if _is_sequence(obj):
 993            return self.transform_point(obj)
 994        obj.apply_transform(self)
 995        return obj
 996
 997########################################################################
 998class TransformInterpolator:
 999    """
1000    Interpolate between a set of linear transformations.
1001    
1002    Position, scale and orientation (i.e., rotations) are interpolated separately,
1003    and can be interpolated linearly or with a spline function.
1004    Note that orientation is interpolated using quaternions via
1005    SLERP (spherical linear interpolation) or the special `vtkQuaternionSpline` class.
1006
1007    To use this class, add at least two pairs of (t, transformation) with the add() method.
1008    Then interpolate the transforms with the `TransformInterpolator(t)` call method,
1009    where "t" must be in the range of `(min, max)` times specified by the add() method.
1010
1011    Example:
1012        ```python
1013        from vedo import *
1014
1015        T0 = LinearTransform()
1016        T1 = LinearTransform().rotate_x(90).shift([12,0,0])
1017
1018        TRI = TransformInterpolator("linear")
1019        TRI.add(0, T0)
1020        TRI.add(1, T1)
1021
1022        plt = Plotter(axes=1)
1023        for i in range(11):
1024            t = i/10
1025            T = TRI(t)
1026            plt += Cube().color(i).apply_transform(T)
1027        plt.show().close()
1028        ```
1029        ![](https://vedo.embl.es/images/other/transf_interp.png)
1030    """
1031    def __init__(self, mode="linear") -> None:
1032        """
1033        Interpolate between two or more linear transformations.
1034        """
1035        self.vtk_interpolator = vtki.new("TransformInterpolator")
1036        self.mode(mode)
1037        self.TS: List[LinearTransform] = []
1038
1039    def __call__(self, t):
1040        """
1041        Get the intermediate transformation at time `t`.
1042        """
1043        xform = vtki.vtkTransform()
1044        self.vtk_interpolator.InterpolateTransform(t, xform)
1045        return LinearTransform(xform)
1046
1047    def add(self, t, T) -> "TransformInterpolator":
1048        """Add intermediate transformations."""
1049        try:
1050            # in case a vedo object is passed
1051            T = T.transform
1052        except AttributeError:
1053            pass
1054        self.TS.append(T)
1055        self.vtk_interpolator.AddTransform(t, T.T)
1056        return self
1057
1058    # def remove(self, t) -> "TransformInterpolator":
1059    #     """Remove intermediate transformations."""
1060    #     self.TS.pop(t)
1061    #     self.vtk_interpolator.RemoveTransform(t)
1062    #     return self
1063    
1064    def trange(self) -> np.ndarray:
1065        """Get interpolation range."""
1066        tmin = self.vtk_interpolator.GetMinimumT()
1067        tmax = self.vtk_interpolator.GetMaximumT()
1068        return np.array([tmin, tmax])
1069    
1070    def clear(self) -> "TransformInterpolator":
1071        """Clear all intermediate transformations."""
1072        self.TS = []
1073        self.vtk_interpolator.Initialize()
1074        return self
1075    
1076    def mode(self, m) -> "TransformInterpolator":
1077        """Set interpolation mode ('linear' or 'spline')."""
1078        if m == "linear":
1079            self.vtk_interpolator.SetInterpolationTypeToLinear()
1080        elif m == "spline":
1081            self.vtk_interpolator.SetInterpolationTypeToSpline()
1082        else:
1083            print('In TransformInterpolator mode can be either "linear" or "spline"')
1084        return self
1085    
1086    @property
1087    def ntransforms(self) -> int:
1088        """Get number of transformations."""
1089        return self.vtk_interpolator.GetNumberOfTransforms()
1090
1091
1092########################################################################
1093# 2d ######
1094def cart2pol(x, y) -> np.ndarray:
1095    """2D Cartesian to Polar coordinates conversion."""
1096    theta = np.arctan2(y, x)
1097    rho = np.hypot(x, y)
1098    return np.array([rho, theta])
1099
1100
1101def pol2cart(rho, theta) -> np.ndarray:
1102    """2D Polar to Cartesian coordinates conversion."""
1103    x = rho * np.cos(theta)
1104    y = rho * np.sin(theta)
1105    return np.array([x, y])
1106
1107
1108########################################################################
1109# 3d ######
1110def cart2spher(x, y, z) -> np.ndarray:
1111    """3D Cartesian to Spherical coordinate conversion."""
1112    hxy = np.hypot(x, y)
1113    rho = np.hypot(hxy, z)
1114    theta = np.arctan2(hxy, z)
1115    phi = np.arctan2(y, x)
1116    return np.array([rho, theta, phi])
1117
1118
1119def spher2cart(rho, theta, phi) -> np.ndarray:
1120    """3D Spherical to Cartesian coordinate conversion."""
1121    st = np.sin(theta)
1122    sp = np.sin(phi)
1123    ct = np.cos(theta)
1124    cp = np.cos(phi)
1125    rst = rho * st
1126    x = rst * cp
1127    y = rst * sp
1128    z = rho * ct
1129    return np.array([x, y, z])
1130
1131
1132def cart2cyl(x, y, z) -> np.ndarray:
1133    """3D Cartesian to Cylindrical coordinate conversion."""
1134    rho = np.sqrt(x * x + y * y)
1135    theta = np.arctan2(y, x)
1136    return np.array([rho, theta, z])
1137
1138
1139def cyl2cart(rho, theta, z) -> np.ndarray:
1140    """3D Cylindrical to Cartesian coordinate conversion."""
1141    x = rho * np.cos(theta)
1142    y = rho * np.sin(theta)
1143    return np.array([x, y, z])
1144
1145
1146def cyl2spher(rho, theta, z) -> np.ndarray:
1147    """3D Cylindrical to Spherical coordinate conversion."""
1148    rhos = np.sqrt(rho * rho + z * z)
1149    phi = np.arctan2(rho, z)
1150    return np.array([rhos, phi, theta])
1151
1152
1153def spher2cyl(rho, theta, phi) -> np.ndarray:
1154    """3D Spherical to Cylindrical coordinate conversion."""
1155    rhoc = rho * np.sin(theta)
1156    z = rho * np.cos(theta)
1157    return np.array([rhoc, phi, z])