├── .github └── workflows │ └── publish.yaml ├── .gitignore ├── LICENSE ├── MANIFEST.in ├── MANYLINUX.md ├── README.md ├── lie_learn ├── __init__.py ├── broadcasting.py ├── groups │ ├── SE2.py │ ├── SO2.py │ ├── SO3.pyx │ ├── SO3_tests.py │ └── __init__.py ├── probability │ ├── HarmonicDensity.py │ ├── S1HarmonicDensity.py │ ├── S2HarmonicDensity.py │ ├── SO3HarmonicDensity.py │ └── __init__.py ├── representations │ ├── SO3 │ │ ├── __init__.py │ │ ├── clebsch_gordan_numerical.py │ │ ├── indexing.py │ │ ├── irrep_bases.pyx │ │ ├── pinchon_hoggan │ │ │ ├── J_block_0-150.npy │ │ │ ├── J_dense_0-150.npy │ │ │ ├── __init__.py │ │ │ ├── download.py │ │ │ ├── pinchon_hoggan.pyx │ │ │ ├── pinchon_hoggan_dense.py │ │ │ └── pinchon_hoggan_parsing.py │ │ ├── spherical_harmonics.py │ │ ├── test_SO3_irrep_bases.py │ │ ├── test_spherical_harmonics.py │ │ ├── test_wigner_d.py │ │ └── wigner_d.py │ └── __init__.py ├── spaces │ ├── S2.py │ ├── S3.py │ ├── Tn.py │ ├── __init__.py │ ├── rn.py │ └── spherical_quadrature.pyx └── spectral │ ├── FFTBase.py │ ├── PolarFFT.py │ ├── S2FFT.py │ ├── S2FFT_NFFT.py │ ├── S2_conv.py │ ├── SE2FFT.py │ ├── SO3FFT_Naive.py │ ├── SO3_conv.py │ ├── T1FFT.py │ ├── T2FFT.py │ ├── __init__.py │ └── fourier_interpolation.py ├── pyproject.toml ├── setup.py └── tests ├── __init__.py ├── spaces ├── __init__.py ├── test_S3_quadrature.py └── test_spherical_quadrature.py └── spectral ├── __init__.py ├── test_S2FFT.py ├── test_S2FFT_NFFT.py ├── test_S2_conv.py ├── test_SO3_FFT_Naive.py └── test_conv_S2_SO3.py /.github/workflows/publish.yaml: -------------------------------------------------------------------------------- 1 | name: "Build and Publish" 2 | 3 | on: 4 | release: 5 | types: [published] 6 | workflow_dispatch: 7 | 8 | 9 | jobs: 10 | build_wheels: 11 | name: Build wheels on ${{ matrix.os }} 12 | runs-on: ${{ matrix.os }} 13 | environment: release 14 | permissions: 15 | id-token: write 16 | strategy: 17 | matrix: 18 | os: [ubuntu-latest, windows-latest, macos-13, macos-14] 19 | 20 | steps: 21 | - uses: actions/checkout@v4 22 | 23 | # Used to host cibuildwheel 24 | - uses: actions/setup-python@v5 25 | 26 | - name: Install cibuildwheel 27 | run: python -m pip install cibuildwheel 28 | 29 | - name: Build wheels 30 | run: python -m cibuildwheel --output-dir dist 31 | # to supply options, put them in 'env', like: 32 | env: 33 | CIBW_BUILD: "cp39-* cp310-* cp311-* cp312-*" 34 | CIBW_SKIP: "*_i686 *-musllinux_*" 35 | - uses: actions/upload-artifact@v4 36 | with: 37 | name: lie_learn-${{ matrix.os }}-${{ strategy.job-index }} 38 | path: ./dist/*.whl 39 | 40 | - name: Publish package 41 | if: matrix.os == 'ubuntu-latest' 42 | uses: pypa/gh-action-pypi-publish@release/v1 43 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Current project only: 2 | ####################### 3 | # ignore cython-generated files 4 | *.c 5 | 6 | # Compiled source # 7 | ################### 8 | *.com 9 | *.class 10 | *.dll 11 | *.exe 12 | *.o 13 | *.so 14 | *.pyc 15 | 16 | # LaTeX # 17 | ######### 18 | *.aux 19 | *.glo 20 | *.idx 21 | *.log 22 | *.toc 23 | *.ist 24 | *.acn 25 | *.acr 26 | *.alg 27 | *.bbl 28 | *.blg 29 | *.dvi 30 | *.glg 31 | *.gls 32 | *.ilg 33 | *.ind 34 | *.lof 35 | *.lot 36 | *.maf 37 | *.mtc 38 | *.mtc1 39 | *.out 40 | *.synctex.gz 41 | 42 | # Packages # 43 | ############ 44 | # it's better to unpack these files and commit the raw source 45 | # git has its own built in compression methods 46 | *.7z 47 | *.dmg 48 | *.gz 49 | *.iso 50 | *.jar 51 | *.rar 52 | *.tar 53 | *.zip 54 | 55 | # Logs and databases # 56 | ###################### 57 | *.log 58 | *.sql 59 | *.sqlite 60 | 61 | # OS generated files # 62 | ###################### 63 | .DS_Store 64 | .DS_Store? 65 | ._* 66 | .Spotlight-V100 67 | .Trashes 68 | Icon? 69 | ehthumbs.db 70 | Thumbs.db 71 | 72 | # Emacs # 73 | ######### 74 | *~ 75 | \#*\# 76 | /.emacs.desktop 77 | /.emacs.desktop.lock 78 | .elc 79 | auto-save-list 80 | tramp 81 | .\#* 82 | 83 | # Others # 84 | ########## 85 | # Python-pickled data files 86 | *.pkl 87 | *.npy 88 | *.imc 89 | *.mat 90 | *.idea 91 | 92 | .eggs 93 | *.egg-info 94 | dist 95 | build 96 | *.html 97 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017, Taco Cohen 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy 4 | of this software and associated documentation files (the "Software"), to deal 5 | in the Software without restriction, including without limitation the rights 6 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 7 | copies of the Software, and to permit persons to whom the Software is 8 | furnished to do so, subject to the following conditions: 9 | 10 | The above copyright notice and this permission notice shall be included in all 11 | copies or substantial portions of the Software. 12 | 13 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 14 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 15 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 16 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 17 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 18 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 19 | SOFTWARE. -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | global-exclude tests/* 2 | global-exclude test_*.py 3 | -------------------------------------------------------------------------------- /MANYLINUX.md: -------------------------------------------------------------------------------- 1 | [Install docker](https://docs.docker.com/get-docker/). 2 | 3 | Get the [manylinux docker environment](https://github.com/pypa/manylinux). 4 | At the time of this writing, `manylinux1` is compatible; however, I used `manylinx2014`. 5 | There is a tool which will label the manylinux binary with the oldest compatible standard. 6 | 7 | ```bash 8 | docker pull quay.io/pypa/manylinux2014_x86_64 9 | ``` 10 | 11 | Run an interactive bash shell in the manylinux docker environment. 12 | 13 | ```bash 14 | docker run -it quay.io/pypa/manylinux2014_x86_64 /bin/bash 15 | ``` 16 | 17 | Inside the interactive bash shell for the docker environment, download lie_learn and change to the source directory. 18 | 19 | ```bash 20 | git clone https://github.com/AMLab-Amsterdam/lie_learn.git 21 | cd lie_learn 22 | ``` 23 | 24 | Create wheels. You have to determine which versions of python are appropriate. 25 | 26 | ```bash 27 | /opt/python/cp35-cp35m/bin/python setup.py bdist_wheel 28 | /opt/python/cp36-cp36m/bin/python setup.py bdist_wheel 29 | /opt/python/cp37-cp37m/bin/python setup.py bdist_wheel 30 | /opt/python/cp38-cp38/bin/python setup.py bdist_wheel 31 | ``` 32 | 33 | Use auditwheel to check for success and modify the binaries to be labeled with the oldest compatible standard (lowest 34 | priority). 35 | 36 | ```bash 37 | auditwheel repair ./dist/lie_learn-0.0.1.post1-cp35-cp35m-linux_x86_64.whl -w ./manylinux 38 | auditwheel repair ./dist/lie_learn-0.0.1.post1-cp36-cp36m-linux_x86_64.whl -w ./manylinux 39 | auditwheel repair ./dist/lie_learn-0.0.1.post1-cp37-cp37m-linux_x86_64.whl -w ./manylinux 40 | auditwheel repair ./dist/lie_learn-0.0.1.post1-cp38-cp38-linux_x86_64.whl -w ./manylinux 41 | ``` 42 | 43 | Open a new terminal window (host environment) and get the running docker `CONTAINER ID`. 44 | 45 | ```bash 46 | docker ps 47 | ``` 48 | 49 | yields 50 | 51 | ``` 52 | CONTAINER ID IMAGE COMMAND CREATED STATUS PORTS NAMES 53 | 8e2b2c3baa8e quay.io/pypa/manylinux2014_x86_64 "/bin/bash" 30 minutes ago Up 30 minutes charming_shannon 54 | ``` 55 | 56 | In my case, the `CONTAINER ID` is `8e2b2c3baa8e`. 57 | In the new terminal window, copy the manylinux wheels from the running container to a folder you'll remember. 58 | 59 | ```bash 60 | mkdir ~/manylinux 61 | docker cp 8e2b2c3baa8e:/lie_learn/manylinux/lie_learn-0.0.1.post1-cp35-cp35m-manylinux1_x86_64.whl ~/manylinux/ 62 | docker cp 8e2b2c3baa8e:/lie_learn/manylinux/lie_learn-0.0.1.post1-cp36-cp36m-manylinux1_x86_64.whl ~/manylinux/ 63 | docker cp 8e2b2c3baa8e:/lie_learn/manylinux/lie_learn-0.0.1.post1-cp37-cp37m-manylinux1_x86_64.whl ~/manylinux/ 64 | docker cp 8e2b2c3baa8e:/lie_learn/manylinux/lie_learn-0.0.1.post1-cp38-cp38-manylinux1_x86_64.whl ~/manylinux/ 65 | ``` 66 | 67 | First do a test by uploading to test pypi. 68 | 69 | ```bash 70 | twine upload --repository-url https://test.pypi.org/legacy/ ~/manylinux/* 71 | ``` 72 | 73 | Try downloading and testing `lie_learn` from there before proceeding. 74 | This is easier said than done. You will need to download all of the dependencies manually then download from test 75 | pypi without any dependencies using 76 | `pip install --no-cache-dir --index-url https://test.pypi.org/simple/ --no-deps lie_learn`. 77 | 78 | Once you know it's working, upload the wheels to pypi with twine. 79 | 80 | ```bash 81 | twine upload ~/manylinux/* 82 | ``` 83 | 84 | For a bit more info, another useful resource is https://opensource.com/article/19/2/manylinux-python-wheels. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | lie_learn is a python package that knows how to do various tricky computations related to Lie groups and manifolds (mainly the sphere S2 and rotation group SO3). This package was written to support various machine learning projects, such as Harmonic Exponential Families [2], (continuous) Group Equivariant Networks [3], Steerable CNNs [4] and Spherical CNNs [5]. 2 | 3 | # What this code can do 4 | - Reparamterize rotations, e.g. matrix to Euler angles to quaternions, etc. (see groups & spaces modules) 5 | - Compute the Wigner-d and Wigner-D matrices (the irreducible representations of SO(3)), and spherical harmonics, using the method developed by Pinchon & Hoggan [1] (see pinchon_hoggan_dense.py). This is a very fast and stable method, but requires a fairly large "J matrix", which we have precomputed up to order 278 using a Maple script. The code will automatically download it from Google Drive during installation. 6 | Note: There are many normalization and phase conventions for both the real and complex versions of the D-matrices and spherical harmonics, and the code can convert between a lot of them (irrep_bases.pyx). 7 | - Compute generalized / non-commutative FFTs for the sphere S2, rotation group SO3, and special Euclidean group SE2 (see spectral module). 8 | - Fit Harmonic Exponential Families on the sphere (probability module; not sure code is still working) 9 | 10 | # Installation 11 | lie_learn can be installed from pypi using: 12 | 13 | ``` 14 | pip install lie_learn 15 | ``` 16 | 17 | Although cython is not a necessary dependency, if you have cython installed, cython will write new versions of the `*.c 18 | ` files before compiling them into `*.so` during installation. To use lie_learn, you will need a c compiler which is 19 | available to python setuptools. 20 | 21 | 22 | # Feedback 23 | For questions and comments, feel free to contact Taco Cohen (http://ta.co.nl). 24 | 25 | 26 | # References 27 | [1] Pinchon, D., & Hoggan, P. E. (2007). Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes. Journal of Physics A: Mathematical and Theoretical, 40(7), 1597–1610. 28 | 29 | [2] Cohen, T. S., & Welling, M. (2015). Harmonic Exponential Families on Manifolds. In Proceedings of the 32nd International Conference on Machine Learning (ICML) (pp. 1757–1765). 30 | 31 | [3] Cohen, T. S., & Welling, M. (2016). Group equivariant convolutional networks. In Proceedings of The 33rd International Conference on Machine Learning (ICML) (Vol. 48, pp. 2990–2999). 32 | 33 | [4] Cohen, T. S., & Welling, M. (2017). Steerable CNNs. In ICLR. 34 | 35 | [5] T.S. Cohen, M. Geiger, J. Koehler, M. Welling (2017). Convolutional Networks for Spherical Signals. In ICML Workshop on Principled Approaches to Deep Learning. 36 | -------------------------------------------------------------------------------- /lie_learn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/__init__.py -------------------------------------------------------------------------------- /lie_learn/broadcasting.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from numpy.lib.index_tricks import as_strided 4 | 5 | 6 | def generalized_broadcast(arrays): 7 | """ 8 | Broadcast X and Y, while ignoring the last axis of X and Y. 9 | 10 | If X.shape = xs + (i,) 11 | and Y.shape = ys + (j,) 12 | then the output arrays have shapes 13 | Xb.shape = zs + (i,) 14 | Yb.shape = zs + (j,) 15 | where zs is the shape of the broadcasting of xs and ys shaped arrays. 16 | 17 | :param arrays: a list of numpy arrays to be broadcasted while ignoring the last axis. 18 | :return: a list of arrays whose shapes have been broadcast 19 | """ 20 | arrays1 = np.broadcast_arrays(*[A[..., 0] for A in arrays]) 21 | shapes_b = [A1.shape + (A.shape[-1],) for A1, A in zip(arrays1, arrays)] 22 | strides_b = [A1.strides + (A.strides[-1],) for A1, A in zip(arrays1, arrays)] 23 | arrays_b = [as_strided(A, shape=shape_Ab, strides=strides_Ab) 24 | for A, shape_Ab, strides_Ab in zip(arrays, shapes_b, strides_b)] 25 | return arrays_b 26 | 27 | 28 | def make_gufunc(f, core_dims_in, core_dims_out): 29 | """ 30 | Automatically turn a function f into a generalized universal function (gufunc). 31 | 32 | :param f: 33 | :param core_dims_in: 34 | :param core_dims_out: 35 | :return: 36 | """ 37 | 38 | return 39 | 40 | def gufunc(args): 41 | args = generalized_broadcast(args) 42 | data_shape = args[0].shape[:-len(core_dims_in[0])] 43 | args = [A.reshape(-1, A.shape[-1]) for A in args] 44 | 45 | #if X_out is None: 46 | # X_out = np.empty_like(X) 47 | #X_out = X_out.reshape(-1, X.shape[-1]) 48 | 49 | out = f(args) 50 | 51 | return out.reshape() 52 | 53 | return gufunc 54 | -------------------------------------------------------------------------------- /lie_learn/groups/SE2.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | import spaces.Rn as Rn 5 | 6 | parameterizations = ('rotation-translation', '2x3 matrix', '3x3 matrix') 7 | 8 | 9 | def compose(g, h, parameterization=None, g_parameterization=None, h_parameterization=None, out_parameterization=None): 10 | """ 11 | Compose elements g, h in SE(2). 12 | """ 13 | if parameterization is not None: 14 | g_parameterization = parameterization 15 | h_parameterization = parameterization 16 | out_parameterization = parameterization 17 | 18 | g_mat = change_parameterization(g, p_from=g_parameterization, p_to='3x3 matrix') 19 | h_mat = change_parameterization(h, p_from=h_parameterization, p_to='3x3 matrix') 20 | gh_mat = np.einsum('...ij,...jk->...ik', g_mat, h_mat) 21 | return change_parameterization(g=gh_mat, p_from='3x3 matrix', p_to=out_parameterization) 22 | 23 | def invert(g, parameterization): 24 | """ 25 | Invert element g in SE(2), where g can have any supported parameterization. 26 | """ 27 | 28 | # Change to (theta, tau1, tau2) paramterization. 29 | g_rt = change_parameterization(g, p_from=parameterization, p_to='rotation-translation') 30 | g_inv_rt = np.empty_like(g_rt) 31 | g_inv_rt[..., 0] = -g_rt[..., 0] 32 | g_inv_rt[..., 1] = -(np.cos(-g_rt[..., 0]) * g_rt[..., 1] - np.sin(-g_rt[..., 0]) * g_rt[..., 2]) 33 | g_inv_rt[..., 2] = -(np.sin(-g_rt[..., 0]) * g_rt[..., 1] + np.cos(-g_rt[..., 0]) * g_rt[..., 2]) 34 | 35 | return change_parameterization(g=g_inv_rt, p_from='rotation-translation', p_to=parameterization) 36 | 37 | 38 | def transform(g, g_parameterization, x, x_parameterization): 39 | """ 40 | Apply rotation g in SE(2) to points x. 41 | """ 42 | g_3x3 = change_parameterization(g, p_from=g_parameterization, p_to='3x3 matrix') 43 | x_homvec = Rn.change_coordinates(x, n=2, p_from=x_parameterization, p_to='homogeneous') 44 | #gx_homvec = g_3x3.dot(x_homvec) 45 | gx_homvec = np.einsum('...ij,...j->...i', g_3x3, x_homvec) 46 | return Rn.change_coordinates(gx_homvec, n=2, p_from='homogeneous', p_to=x_parameterization) 47 | 48 | 49 | def change_parameterization(g, p_from, p_to): 50 | 51 | g = np.array(g) 52 | 53 | if p_from == p_to: 54 | return g 55 | 56 | if p_from == 'rotation-translation' and p_to == '2x3 matrix': 57 | g_out = np.empty(g.shape[:-1] + (2, 3)) 58 | g_out[..., 0, 0] = np.cos(g[..., 0]) 59 | g_out[..., 0, 1] = -np.sin(g[..., 0]) 60 | g_out[..., 1, 0] = np.sin(g[..., 0]) 61 | g_out[..., 1, 1] = np.cos(g[..., 0]) 62 | g_out[..., 0, 2] = g[..., 1] 63 | g_out[..., 1, 2] = g[..., 2] 64 | return g_out 65 | 66 | if p_from == 'rotation-translation' and p_to == '3x3 matrix': 67 | g_out = np.empty(g.shape[:-1] + (3, 3)) 68 | g_out[..., 0, 0] = np.cos(g[..., 0]) 69 | g_out[..., 0, 1] = -np.sin(g[..., 0]) 70 | g_out[..., 0, 2] = g[..., 1] 71 | g_out[..., 1, 0] = np.sin(g[..., 0]) 72 | g_out[..., 1, 1] = np.cos(g[..., 0]) 73 | g_out[..., 1, 2] = g[..., 2] 74 | g_out[..., 2, 0] = 0. 75 | g_out[..., 2, 1] = 0. 76 | g_out[..., 2, 2] = 1. 77 | return g_out 78 | 79 | else: 80 | raise ValueError('Not supported (yet):' + str(p_from) + ' to ' + str(p_to)) -------------------------------------------------------------------------------- /lie_learn/groups/SO2.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | parameterizations = ('MAT', 'C', 'ANG') 5 | 6 | def compose(g, h, parameterization='MAT', g_parameterization=None, h_parameterization=None, out_parameterization=None): 7 | """ 8 | Compose elements g, h in SO(2). 9 | g and h can have the following parameterizations: 10 | 1: MAT 2x2 rotation matrix 11 | 2: C 1x1 complex exponential (z=exp(i theta)) 12 | """ 13 | if parameterization is not None: 14 | g_parameterization = parameterization 15 | h_parameterization = parameterization 16 | out_parameterization = parameterization 17 | 18 | g_mat = change_parameterization(g, p_from=g_parameterization, p_to='MAT') 19 | h_mat = change_parameterization(h, p_from=h_parameterization, p_to='MAT') 20 | gh_mat = np.einsum('...ij,...jk->...ik', g_mat, h_mat) 21 | return change_parameterization(g=gh_mat, p_from='MAT', p_to=out_parameterization) 22 | 23 | def invert(g, parameterization): 24 | """ 25 | Invert element g in SO(2), where g can have any supported parameterization: 26 | 1: MAT 2x2 rotation matrix 27 | 2: C 1x1 complex exponential (z=exp(i theta)) 28 | """ 29 | 30 | g_mat = change_parameterization(g, p_from=parameterization, p_to='MAT') 31 | g_mat_T = g_mat.transpose(list(range(0, g_mat.ndim - 2)) + [g_mat.ndim - 1, g_mat.ndim - 2]) # Transpose last axes 32 | return change_parameterization(g=g_mat_T, p_from='MAT', p_to=parameterization) 33 | 34 | def transform(g, g_parameterization, x, x_parameterization): 35 | """ 36 | Apply rotation g in SO(2) to points x. 37 | """ 38 | #g_mat = change_parameterization(g_parameterization + 'toMAT', g, ichk=0) 39 | #x_vec = change_coordinates(x_parameterization + 'toC', x) 40 | #gx_vec = g_mat.dot(x_vec) 41 | #return change_coordinates('Cto' + x_parameterization, gx_vec) 42 | raise NotImplementedError('SO2 transform not implemented') 43 | 44 | 45 | def change_parameterization(g, p_from, p_to): 46 | """ 47 | 48 | """ 49 | if p_from == p_to: 50 | return g 51 | 52 | elif p_from == 'MAT' and p_to == 'C': 53 | theta = np.arctan2(g[..., 1, 0], g[..., 0, 0]) 54 | return np.exp(1j * theta) 55 | elif p_from == 'MAT' and p_to == 'ANG': 56 | return np.arctan2(g[..., 1, 0], g[..., 0, 0]) 57 | elif p_from == 'C' and p_to == 'MAT': 58 | theta = np.angle(g) 59 | c = np.cos(theta) 60 | s = np.sin(theta) 61 | return np.array([[c, -s], [s, c]]).transpose(list(range(2, 2 + c.ndim)) + [0, 1]) 62 | elif p_from == 'C' and p_to == 'ANG': 63 | return np.angle(g) 64 | elif p_from == 'ANG' and p_to == 'MAT': 65 | c = np.cos(g) 66 | s = np.sin(g) 67 | return np.array([[c, -s], [s, c]]).transpose(list(range(2, 2 + c.ndim)) + [0, 1]) 68 | elif p_from == 'ANG' and p_to == 'C': 69 | return np.exp(1j * g) 70 | else: 71 | raise ValueError('Unsupported conversion:' + str(p_from) + ' to ' + str(p_to)) -------------------------------------------------------------------------------- /lie_learn/groups/SO3_tests.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from lie_learn.groups.SO3 import * 4 | 5 | 6 | def test_change_parameterization(): 7 | 8 | def is_equal(R1, R2, p): 9 | if p == 'Q': 10 | # Quaternions are only defined up to a sign, so check each row, what sign we need 11 | for i in range(R1.shape[0]): 12 | if not (np.allclose(R1[i, ...], R2[i, ...]) or np.allclose(R1[i, ...], -R2[i, ...])): 13 | return False 14 | return True 15 | elif p == 'EV': 16 | # Euler vector (x,y,z,theta) == (-x,-y,-z,-theta mod 2pi) 17 | for i in range(R1.shape[0]): 18 | R2i = np.array([-R2[i, 0], -R2[i, 1], -R2[i, 2], (-R2[i, 3]) % (2 * np.pi)]) 19 | if not (np.allclose(R1[i, ...], R2[i, ...]) or np.allclose(R1[i, ], R2i)): 20 | return False 21 | return True 22 | 23 | else: 24 | return np.allclose(R1, R2) 25 | 26 | for p1 in parameterizations: 27 | for p2 in parameterizations: 28 | 29 | # Create two random rotations in 313 Euler angles 30 | R1_EA313 = (np.random.rand(3) * np.array([2 * np.pi, np.pi, 2 * np.pi]))[np.newaxis, :] 31 | R2_EA313 = (np.random.rand(3) * np.array([2 * np.pi, np.pi, 2 * np.pi]))[np.newaxis, :] 32 | R_EA313 = np.r_[R1_EA313, R2_EA313] 33 | 34 | R1_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R1_EA313) 35 | R1_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R1_EA313) 36 | R2_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R2_EA313) 37 | R2_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R2_EA313) 38 | R_p1 = change_coordinates(p_from='EA313', p_to=p1, g=R_EA313) 39 | R_p2 = change_coordinates(p_from='EA313', p_to=p2, g=R_EA313) 40 | 41 | R1_p2_from_R1_p1 = change_coordinates(p_from=p1, p_to=p2, g=R1_p1) 42 | R1_p1_from_R1_p2 = change_coordinates(p_from=p2, p_to=p1, g=R1_p2) 43 | R2_p2_from_R2_p1 = change_coordinates(p_from=p1, p_to=p2, g=R2_p1) 44 | R2_p1_from_R2_p2 = change_coordinates(p_from=p2, p_to=p1, g=R2_p2) 45 | R_p2_from_R_p1 = change_coordinates(p_from=p1, p_to=p2, g=R_p1) 46 | R_p1_from_R_p2 = change_coordinates(p_from=p2, p_to=p1, g=R_p2) 47 | 48 | assert is_equal(R1_p1_from_R1_p2, R1_p1, p1), ( 49 | p1 + ' to ' + p2 + ' | R1_p1: ' + str(R1_p1) + ' | R1_p2: ' + str(R1_p2) + ' | R1_p1_from_R1_p2: ' + 50 | str(R1_p1_from_R1_p2)) 51 | assert is_equal(R2_p1_from_R2_p2, R2_p1, p1), ( 52 | p1 + ' to ' + p2 + ' | R2_p1: ' + str(R2_p1) + ' | R2_p2: ' + str(R2_p2) + ' | R2_p1_from_R2_p2: ' + 53 | str(R2_p1_from_R2_p2)) 54 | assert is_equal(R_p1_from_R_p2, R_p1, p1), ( 55 | p1 + ' to ' + p2 + ' | R_p1: ' + str(R_p1) + ' | R_p2: ' + str(R_p2) + ' | R_p1_from_R_p2: ' + 56 | str(R_p1_from_R_p2)) 57 | assert is_equal(R1_p2_from_R1_p1, R1_p2, p2), ( 58 | p1 + ' to ' + p2 + ' | R1_p1: ' + str(R1_p1) + ' | R1_p2: ' + str(R1_p2) + ' | R1_p2_from_R1_p1: ' + 59 | str(R1_p2_from_R1_p1)) 60 | assert is_equal(R2_p2_from_R2_p1, R2_p2, p2), ( 61 | p1 + ' to ' + p2 + ' | R2_p1: ' + str(R2_p1) + ' | R2_p2: ' + str(R2_p2) + ' | R2_p2_from_R2_p1: ' + 62 | str(R2_p2_from_R2_p1)) 63 | assert is_equal(R_p2_from_R_p1, R_p2, p2), ( 64 | p1 + ' to ' + p2 + ' | R_p1: ' + str(R_p1) + ' | R_p2: ' + str(R_p2) + ' | R_p2_from_R_p1: ' + 65 | str(R_p2_from_R_p1)) 66 | 67 | def test_invert(): 68 | 69 | for p in parameterizations: 70 | 71 | R_EA = np.random.rand(4, 5, 6, 3) * np.array([2 * np.pi, np.pi, 2 * np.pi])[None, None, None, :] 72 | R_p = change_coordinates(R_EA, p_from='EA313', p_to=p) 73 | R_p_inv = invert(R_p, parameterization=p) 74 | 75 | e = compose(R_p, R_p_inv, parameterization=p) 76 | eM = change_coordinates(e, p_from=p, p_to='MAT') 77 | assert np.isclose(np.sum(eM - np.eye(3)), 0.0), 'not the identity: ' + eM -------------------------------------------------------------------------------- /lie_learn/groups/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/groups/__init__.py -------------------------------------------------------------------------------- /lie_learn/probability/HarmonicDensity.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | class HarmonicDensity(object): 4 | 5 | pass 6 | -------------------------------------------------------------------------------- /lie_learn/probability/S1HarmonicDensity.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | import numpy as np 4 | from scipy.fftpack import rfft, irfft 5 | from scipy.optimize import fmin_l_bfgs_b 6 | 7 | """ 8 | samples = vonmises.rvs(kappa=1.0, size=10000) 9 | J = 10; alpha = 10; n = (2 * J + 1) * alpha 10 | even = np.arange(0, 2 * J, 2); odd = np.arange(1, 2 * J, 2) 11 | empirical_moments = np.zeros(2 * J) 12 | empirical_moments[even] = np.mean(np.cos(np.arange(1, J + 1)[np.newaxis, :] * samples[:, np.newaxis]), axis=0) 13 | empirical_moments[odd] = np.mean(np.sin(np.arange(1, J + 1)[np.newaxis, :] * samples[:, np.newaxis]), axis=0) 14 | 15 | 16 | def logp_and_grad(eta): 17 | # Compute moments: 18 | negative_energy = irfft(np.hstack([[0], eta]), n=n) * (n / 2.) 19 | unnormalized_moments = rfft(np.exp(negative_energy)) / (n / 2.) * np.pi 20 | Z = unnormalized_moments[0] 21 | moments = unnormalized_moments[1:eta.size + 1] / Z 22 | 23 | # Compute gradients and log-prob: 24 | grad_logp = empirical_moments - moments 25 | logp = eta.dot(empirical_moments) - np.log(Z) 26 | return -logp, -grad_logp 27 | 28 | opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(logp_and_grad, x0=np.zeros(2 * J), iprint=0, factr=1e7, pgtol=1e-5) 29 | print info['task'] 30 | print 'Optimum log-likelihood:', -opt_neg_logp 31 | print 'Optimal parameters:', np.round(opt_eta, 2) 32 | """ 33 | 34 | 35 | class S2HarmonicDensity(): 36 | 37 | def __init__(self, L_max, oversampling_factor=2): 38 | 39 | self.L_max = L_max 40 | self.L_max_os = self.L_max * oversampling_factor 41 | 42 | self.even = np.arange(0, 2 * L_max, 2) 43 | self.odd = np.arange(1, 2 * L_max, 2) 44 | 45 | self.even_os = np.arange(0, 2 * self.L_max_os, 2) 46 | self.odd_os = np.arange(1, 2 * self.L_max_os, 2) 47 | 48 | 49 | def negative_energy(self, x, eta): 50 | """ 51 | 52 | :param x: 53 | :param eta: 54 | :return: 55 | """ 56 | 57 | pass #return eta.dot(sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis], 58 | # x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :], 59 | # field='real', normalization='quantum', condon_shortley=True)) 60 | 61 | 62 | 63 | def moments(self, eta): 64 | """ 65 | 66 | :param eta: 67 | :return: 68 | """ 69 | pass 70 | 71 | def empirical_moments(self, X, average=True): 72 | """ 73 | Compute the empirical moments of the sample x 74 | 75 | :param x: dataset shape (N, 2) for 2 spherical coordinates (theta, phi) per point 76 | :return: the moments 1/N sum_i=1^N T(x_i) 77 | """ 78 | pass 79 | 80 | def grad_log_p(self, eta, empirical_moments): 81 | """ 82 | 83 | :param eta: 84 | :param M: 85 | :return: 86 | """ 87 | pass 88 | 89 | def log_p_and_grad(self, eta, empirical_moments): 90 | """ 91 | 92 | """ 93 | pass 94 | 95 | 96 | def mle_lbfgs(self, empirical_moments, eta_init=None, SigmaInv=None, verbose=True): 97 | 98 | # Move to base-class? 99 | 100 | if eta_init is None: 101 | eta = np.zeros((self.L_max + 1) ** 2 - 1) 102 | else: 103 | eta = eta_init.copy() 104 | 105 | if SigmaInv is None: # No regularization 106 | def objective_and_grad(eta): 107 | logp, grad = self.log_p_and_grad(eta, empirical_moments) 108 | return -logp, -grad 109 | else: 110 | lnZ_prior = 0.5 * SigmaInv.size * np.log(2 * np.pi) - 0.5 * np.sum(np.log(SigmaInv)) 111 | def objective_and_grad(eta): 112 | logp, grad = self.log_p_and_grad(eta, empirical_moments) 113 | SigmaInv_eta = SigmaInv * eta 114 | logp += -0.5 * eta.dot(SigmaInv_eta) - lnZ_prior 115 | grad += -SigmaInv_eta 116 | return -logp, -grad 117 | 118 | opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(objective_and_grad, x0=eta, iprint=int(verbose) - 1, 119 | factr=1e7, # moderate accuracy 120 | #factr=1e12, # low accuracy 121 | pgtol=1e-5) # norm of proj. grad. to stop iteration at 122 | 123 | if verbose: 124 | print('Maximum log prob:', -opt_neg_logp) 125 | print('Optimization info:', info) 126 | 127 | # Finally, compute Z: 128 | _, lnZ = self.moments(opt_eta) 129 | return opt_eta, lnZ 130 | 131 | 132 | def _moment_numerical_integration(self, eta, l, m): 133 | """ 134 | Compute the (l,m)-moment of the density with natural parameter eta using slow numerical integration. 135 | The output of this function should be equal to the *unnormalized* moment as it comes out of the FFT 136 | (without dividing by Z). 137 | 138 | :param eta: 139 | :param l: 140 | :param m: 141 | :return: 142 | """ 143 | pass 144 | #f = lambda theta, phi: (np.exp(self.negative_energy(np.array([[theta, phi]]), eta)) 145 | # * sh(l, m, theta, phi, 146 | # field='real', normalization='quantum', 147 | # condon_shortley=True)) 148 | #return S2.integrate(f) 149 | 150 | -------------------------------------------------------------------------------- /lie_learn/probability/S2HarmonicDensity.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from scipy.optimize import fmin_cg, fmin_l_bfgs_b 4 | from ..spectral.S2FFT_NFFT import S2FFT_NFFT 5 | from ..spaces import S2 6 | from ..representations.SO3.spherical_harmonics import sh 7 | 8 | 9 | # TODO: add unit-test to check moments agains numerical integration (have done this in terminal) 10 | 11 | class S2HarmonicDensity(): 12 | 13 | def __init__(self, L_max, oversampling_factor=2, fft=None): 14 | 15 | # Compute the maximum degree from the length of eta 16 | # sum_l=0^L (2 l + 1) = (L + 1)^2 17 | # so 18 | # L = sqrt(eta.size) - 1 19 | #if (np.sqrt(eta.shape[-1] + 1) != np.sqrt(eta.shape[-1] + 1).astype(int)).any(): 20 | # raise ValueError('Incorrect eta: last dimension must be a square.') 21 | #self.L_max = int(np.sqrt(eta.shape[-1] + 1) - 1) 22 | self.L_max = L_max 23 | self.L_max_os = self.L_max * oversampling_factor 24 | 25 | # Create arrays containing the (l,m) coordinates 26 | l = [[l] * (2 * l + 1) for l in range(1, self.L_max + 1)] 27 | self.ls = np.array([ll for sublist in l for ll in sublist]) # 1, 1, 1, 2, 2, 2, 2, 2, ... 28 | l_oversampled = [[l] * (2 * l + 1) for l in range(1, self.L_max_os + 1)] 29 | self.ls_oversampled = np.array([ll for sublist in l_oversampled for ll in sublist]) 30 | 31 | m = [list(range(-l, l + 1)) for l in range(1, self.L_max + 1)] 32 | self.ms = np.array([mm for sublist in m for mm in sublist]) # -1, 0, 1, -2, -1, 0, 1, 2, ... 33 | m_oversampled = [list(range(-l, l + 1)) for l in range(1, self.L_max_os + 1)] 34 | self.ms_oversampled = np.array([mm for sublist in m_oversampled for mm in sublist]) 35 | 36 | if fft is None: 37 | # Setup a spherical grid and corresponding quadrature weights 38 | convention = 'Clenshaw-Curtis' 39 | #convention = 'Gauss-Legendre' 40 | 41 | x = S2.meshgrid(b=self.L_max_os, grid_type=convention) 42 | w = S2.quadrature_weights(b=self.L_max_os, grid_type=convention) 43 | self.fft = S2FFT_NFFT(L_max=self.L_max_os, x=x, w=w) 44 | else: 45 | if fft.L_max < self.L_max_os: 46 | raise ValueError('fft.L_max must be larger than or equal to L_max * oversampling_factor') 47 | self.fft = fft 48 | 49 | def negative_energy(self, x, eta): 50 | """ 51 | 52 | :param x: 53 | :param eta: 54 | :return: 55 | """ 56 | x = np.atleast_2d(x) 57 | return eta.dot(sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis], 58 | x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :], 59 | field='real', normalization='quantum', condon_shortley=True)) 60 | 61 | def sufficient_statistics(self, x): 62 | x = np.atleast_2d(x) 63 | return sh(self.ls[:, np.newaxis], self.ms[:, np.newaxis], 64 | x[:, 0][np.newaxis, :], x[:, 1][np.newaxis, :], 65 | field='real', normalization='quantum', condon_shortley=True) 66 | 67 | def moments(self, eta): 68 | """ 69 | 70 | :param eta: 71 | :return: 72 | """ 73 | # 74 | eta_os = np.zeros((self.L_max_os + 1) ** 2) 75 | eta_os[1:eta.size + 1] = eta 76 | 77 | neg_e = self.fft.synthesize(eta_os) 78 | #unnormalized_moments = self.fft.analyze(np.exp(neg_e)) 79 | #return unnormalized_moments[1:] / unnormalized_moments[0], unnormalized_moments[0] 80 | 81 | maximum = np.max(neg_e) 82 | unnormalized_moments = self.fft.analyze(np.exp(neg_e - maximum)) 83 | 84 | #log_unnormalized_moments = np.log(unnormalized_moments + 0j) 85 | #moments = np.exp(log_unnormalized_moments - log_unnormalized_moments[0]).real 86 | #Z = np.exp(log_raw_moments[0] + maximum).real 87 | #lnZ = (log_unnormalized_moments[0] + maximum).real 88 | 89 | unnormalized_moments[0] *= np.sqrt(4 * np.pi) 90 | 91 | moments = unnormalized_moments / unnormalized_moments[0] 92 | #print np.sum(np.abs(m2 - moments)) 93 | lnZ = np.log(unnormalized_moments[0]) + maximum 94 | #print lnZ, lnZ2, lnZ - lnZ2 95 | 96 | return moments[1:(self.L_max + 1) ** 2], lnZ 97 | 98 | def moments_numint(self, eta): 99 | 100 | moments = np.zeros((self.L_max + 1) ** 2) 101 | 102 | f = lambda th, ph: np.exp(self.negative_energy([th, ph], eta)) 103 | moments[0] = S2.integrate(f, normalize=False) 104 | 105 | for l in range(1, self.L_max + 1): 106 | for m in range(-l, l + 1): 107 | print('integrating', l, m) 108 | f = lambda th, ph: np.exp(self.negative_energy([th, ph], eta)) * sh(l, m, th, ph, 109 | field='real', normalization='quantum', condon_shortley=True) 110 | moments[l ** 2 + l + m] = S2.integrate(f, normalize=False) 111 | 112 | return moments[1:] / moments[0], moments[0] 113 | 114 | def empirical_moments(self, X, average=True): 115 | """ 116 | Compute the empirical moments of the sample x 117 | 118 | :param x: dataset shape (N, 2) for 2 spherical coordinates (theta, phi) per point 119 | :return: the moments 1/N sum_i=1^N T(x_i) 120 | """ 121 | # TODO: this can be done potentially more efficiently by computing T(0,0) (the suff. stats. of the north pole), 122 | # and then transforming by D(theta, phi, 0) or something similar. This matrix vector-multiplication can be 123 | # done efficiently by the Pinchon-Hoggan method. (or asymptotically even faster using other methods) 124 | 125 | T = sh(self.ls[np.newaxis, :], self.ms[np.newaxis, :], 126 | X[:, 0][:, np.newaxis], X[:, 1][:, np.newaxis], 127 | field='real', normalization='quantum', condon_shortley=True) 128 | if average: 129 | return T.mean(axis=0) 130 | else: 131 | return T 132 | 133 | def grad_log_p(self, eta, empirical_moments): 134 | """ 135 | 136 | :param eta: 137 | :param M: 138 | :return: 139 | """ 140 | moments, _ = self.moments(eta) 141 | return empirical_moments - moments 142 | 143 | def log_p_and_grad(self, eta, empirical_moments): 144 | """ 145 | Compute the gradient of the log probability of the density given by eta, 146 | evaluated at a sample of data summarized by the empirical moments. 147 | The log-prob is: 148 | ln prod_i=1^N p(x_i | eta) 149 | = 150 | sum_i=1^N eta^T T(x_i) - ln Z_eta 151 | = 152 | N (eta^T T_bar - ln Z_eta) 153 | where T_bar = 1/N sum_i=1^N T(x_i) are the empirical moments, as computed by self.empirical_moments(X). 154 | In this function we work with the *average* log-prob, i.e. leaving out the factor N from the log-prob formula. 155 | 156 | The gradient is (leaving out the factor of N) 157 | T_bar - E_eta[T(x)] 158 | where E_eta[T(x)] are the moments of p(x|eta), as computed by self.moments(eta). 159 | 160 | :param eta: the natural parameters of the distribution 161 | :param empirical_moments: the average sufficient statistics, as computed by self.empirical_moments(X) 162 | :return: the gradient of the average log-prob with respect to eta, and the average log prob itself. 163 | """ 164 | moments, lnZ = self.moments(eta) 165 | grad_logp = empirical_moments - moments 166 | logp = eta.dot(empirical_moments) - lnZ 167 | return logp, grad_logp 168 | 169 | def mle_sgd(self, empirical_moments, eta_init=None, learning_rate=0.1, max_iter=1000, verbose=True): 170 | """ 171 | 172 | :param X: 173 | :return: 174 | """ 175 | if eta_init is None: 176 | eta = np.zeros((self.L_max + 1) ** 2 - 1) 177 | else: 178 | eta = eta_init.copy() 179 | 180 | for i in range(max_iter): 181 | log_p, grad_log_p = self.log_p_and_grad(eta, empirical_moments) 182 | eta += learning_rate * grad_log_p 183 | if verbose: 184 | print('log-prob:', log_p) 185 | 186 | # Finally, compute Z: 187 | _, lnZ = self.moments(eta) 188 | return eta, lnZ 189 | 190 | def mle_lbfgs(self, empirical_moments, eta_init=None, SigmaInv=None, verbose=True): 191 | 192 | if eta_init is None: 193 | eta = np.zeros((self.L_max + 1) ** 2 - 1) 194 | else: 195 | eta = eta_init.copy() 196 | 197 | if SigmaInv is None: # No regularization 198 | def objective_and_grad(eta): 199 | logp, grad = self.log_p_and_grad(eta, empirical_moments) 200 | return -logp, -grad 201 | else: 202 | #lnZ_prior = 0.5 * SigmaInv.size * np.log(2 * np.pi) - 0.5 * np.sum(np.log(SigmaInv)) 203 | def objective_and_grad(eta): 204 | logp, grad = self.log_p_and_grad(eta, empirical_moments) 205 | SigmaInv_eta = SigmaInv * eta 206 | logp += -0.5 * eta.dot(SigmaInv_eta) # - lnZ_prior 207 | grad += -SigmaInv_eta 208 | return -logp, -grad 209 | 210 | opt_eta, opt_neg_logp, info = fmin_l_bfgs_b(objective_and_grad, x0=eta, iprint=int(verbose) - 1, 211 | #factr=1e7, # moderate accuracy 212 | factr=1e12, # low accuracy 213 | pgtol=1e-4, 214 | maxiter=1000) # norm of proj. grad. to stop iteration at 215 | 216 | if verbose: 217 | print('Maximum log prob:', -opt_neg_logp) 218 | print('Optimization info:', info['warnflag'], info['task'], np.mean(info['grad'])) 219 | 220 | # Finally, compute Z: 221 | _, lnZ = self.moments(opt_eta) 222 | return opt_eta, lnZ 223 | 224 | def mle_cg(self, empirical_moments, eta_init=None, verbose=True): 225 | 226 | if eta_init is None: 227 | eta = np.zeros((self.L_max + 1) ** 2 - 1) 228 | else: 229 | eta = eta_init.copy() 230 | 231 | def objective(eta): 232 | logp, _ = self.log_p_and_grad(eta, empirical_moments) 233 | return -logp 234 | def grad(eta): 235 | _, grad = self.log_p_and_grad(eta, empirical_moments) 236 | return -grad 237 | eta_min, logp_min, fun_calls, grad_calls, warnflag = fmin_cg(f=objective, fprime=grad, x0=eta, 238 | full_output=True) 239 | 240 | if verbose: 241 | print('min log p:', logp_min) 242 | print('fun_calls:', fun_calls) 243 | print('grad_calls:', grad_calls) 244 | print('warnflag:', warnflag) 245 | #print 'allvecs:', allvecs 246 | 247 | # Finally, compute Z: 248 | _, lnZ = self.moments(eta_min) 249 | return eta_min, lnZ 250 | 251 | def _moment_numerical_integration(self, eta, l, m): 252 | """ 253 | Compute the (l,m)-moment of the density with natural parameter eta using slow numerical integration. 254 | The output of this function should be equal to the *unnormalized* moment as it comes out of the FFT 255 | (without dividing by Z). 256 | 257 | :param eta: 258 | :param l: 259 | :param m: 260 | :return: 261 | """ 262 | f = lambda theta, phi: (np.exp(self.negative_energy(np.array([[theta, phi]]), eta)) 263 | * sh(l, m, theta, phi, 264 | field='real', normalization='quantum', 265 | condon_shortley=True)) 266 | return S2.integrate(f) 267 | -------------------------------------------------------------------------------- /lie_learn/probability/SO3HarmonicDensity.py: -------------------------------------------------------------------------------- 1 | 2 | from .HarmonicDensity import HarmonicDensity 3 | 4 | class SO3HarmonicDensity(HarmonicDensity): 5 | 6 | def __init__(self): 7 | pass 8 | 9 | 10 | -------------------------------------------------------------------------------- /lie_learn/probability/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/probability/__init__.py -------------------------------------------------------------------------------- /lie_learn/representations/SO3/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/representations/SO3/__init__.py -------------------------------------------------------------------------------- /lie_learn/representations/SO3/clebsch_gordan_numerical.py: -------------------------------------------------------------------------------- 1 | 2 | """ 3 | Compute the Clebsch-Gordan coefficients of SO(3) numerically. 4 | These coefficients specify how a product of Wigner D functions can be re-expressed as a linear combination of 5 | Wigner D functions. 6 | 7 | Wikipedia gives the formula: 8 | D^l_mn(a,b,c) D^l'_m'n'(a,b,c) = sum_{L=|l-l'|^(l+l') sum_M=-L^L sum_N=-L^L D^L_MN(a,b,c) 9 | where D are the Wigner D functions and <....|..> are CG coefficients. 10 | 11 | For our computations related to the representations of SO(3) on the projective plane, we are most interested in 12 | the case l=l'=1 13 | D^1_mn(a,b,c) D^1_m'n'(a,b,c) = sum_{L=0^2 sum_M=-L^L sum_N=-L^L D^L_MN(a,b,c) 14 | 15 | Since we typically employ the Pinchon-Hoggan basis of real spherical harmonics, we cannot use the formulas for the 16 | CG coefficients that are given in the standard references. We could re-do the analysis to find these coefficients 17 | for the real basis, but here we employ a simple numerical approach. 18 | 19 | We view the products coefficients <1m1m'|LM><1n1n'|LN> as unknowns C(m, m', n, n', L, M, N) 20 | We can randomly sample a large number of euler angles g_i = (alpha_i, beta_i, gamma_i), and solve the linear 21 | system for C. 22 | 23 | Looking at this system, it appears that the coefficients are ratios of integers or (square?) roots, such as 24 | 1/2, 1/3, 1/(2 sqrt(3)), etc. 25 | This is reminiscent of the numbers Pinchon & Hoggan encounter in their J matrix, so that is something to look into. 26 | 27 | I've saved dtThe results of the CG computation for l=1 to clebsch_gordan_l1.npy 28 | """ 29 | 30 | import numpy as np 31 | from pinchon_hoggan import * 32 | 33 | 34 | def compute_CG_3D(m1, n1, m2, n2, N=1000): 35 | 36 | l = 1 37 | 38 | l_min = 0 39 | l_max = 2 40 | num_coefs = sum([(2 * j + 1) ** 2 for j in range(l_min, l_max + 1)]) 41 | 42 | g = np.random.rand(3, N) * np.array([[2 * np.pi], [np.pi], [2 * np.pi]]) 43 | 44 | D1 = SO3_irrep(g, 1)[l + m1, l + n1, :] 45 | D2 = SO3_irrep(g, 1)[l + m2, l + n2, :] 46 | target = D1 * D2 47 | 48 | A = np.zeros((N, num_coefs)) 49 | for i in range(N): 50 | Ds = np.concatenate([SO3_irrep(g[:, i][:, None], j).flatten() for j in range(l_min, l_max + 1)]) 51 | A[i, :] = Ds 52 | 53 | return A, target, np.linalg.pinv(A).dot(target) 54 | 55 | 56 | def compute_CG_matrix(N=1000): 57 | 58 | CG = np.zeros((1 + 3 * 3 + 5 * 5, 3, 3, 3, 3)) 59 | l = 1 60 | for m1 in range(-l, l + 1): 61 | for n1 in range(-l, l + 1): 62 | for m2 in range(-l, l + 1): 63 | for n2 in range(-l, l + 1): 64 | print(m1, n1, m2, n2) 65 | _, _, w = compute_CG_3D(m1, n1, m2, n2, N) 66 | CG[:, l + m1, l + n1, l + m2, l + n2] = w 67 | 68 | return CG 69 | 70 | if __name__ == '__main__': 71 | CG = compute_CG_matrix(1000) 72 | CG_exact = np.zeros_like(CG) 73 | 74 | uniques = [0., 1. / 2., -1. / 2., 75 | 1. / 3., -1. / 3., 76 | 1. / 6., 2. / 3., 77 | 1. / (2 * np.sqrt(3)), -1. / (2 * np.sqrt(3)), 78 | 1. / np.sqrt(3), -1. / np.sqrt(3)] 79 | print('Hypothetical exact uniques:') 80 | print(np.sort(uniques)) 81 | print('Numerically obtained uniques (rounded to 5 decimals)') 82 | print(np.unique(np.round(CG, 5))) 83 | for value in uniques: 84 | inds = np.nonzero(np.isclose(CG, value)) 85 | CG_exact[inds] = value 86 | 87 | print('Absolute error between exact and numerical:', np.sum(np.abs(CG_exact - CG))) -------------------------------------------------------------------------------- /lie_learn/representations/SO3/indexing.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def flat_ind_so3(l, m, n): 5 | """ 6 | The SO3 spectrum consists of matrices f_hat^l of size (2l+1, 2l+1) for l=0, ..., L_max. 7 | If we flatten these matrices and stack them, we get a big vector where the element f_hat^l_mn has a certain 8 | flat index. This function computes that index. 9 | 10 | The number of elements up to and including order L is 11 | N_L = sum_{l=0}^L (2l+1)^2 = 1/3 (2 L + 1) (2 L + 3) (L + 1) 12 | 13 | Element (l, m, n) has N_l elements before it from previous blocks, and in addition several elements in the current 14 | block. The number of elements in the current block, before (m, n) is determined as follows. 15 | First we associate with indices m and n (running from -l to l) with their zero-based index: 16 | m' = m + l 17 | n' = n + l 18 | The linear index of this pair (m', n') is 19 | i = m' * w + n' 20 | where w is the width of the matrix, i.e. w = 2l + 1 21 | 22 | The final index of (l, m, n) is N_L + i 23 | 24 | :param l, m, n: spectral indices 25 | :return: index of (l, m, n) in flat vector 26 | """ 27 | assert np.abs(m) <= l 28 | assert np.abs(n) <= l 29 | 30 | if l == 0: 31 | return 0 # The N_L formula only works for l > 0, so we special case this 32 | 33 | L = l - 1 34 | N_L = ((2 * L + 1) * (2 * L + 3) * (L + 1)) // 3 35 | i = (m + l) * (2 * l + 1) + (n + l) 36 | return N_L + i 37 | 38 | 39 | def flat_ind_zp_so3(l, m, n, b): 40 | """ 41 | The SO3 spectrum consists of matrices f_hat^l of size (2l+1, 2l+1) for l=0, ..., L_max = b - 1. 42 | These can be stored in a zero-padded array A of shape (b, 2b, 2b) with axes l, m, n with zero padding around 43 | the center of the last two axes. If we flatten this array A we get a vector v of size 4b^3. 44 | This function gives the flat index in this array v corresponding to element (l, m, n) 45 | 46 | The zero-based 3D index of (l, m, n) in A is (l, b + m, b + n). 47 | The corresponding flat index is i = l * 4b^2 + (b + m) * 2b + b + n 48 | 49 | :param l, m, n: spectral indices 50 | :return: index of (l, m, n) in flat zero-padded vector 51 | """ 52 | return l * 4 * (b ** 2) + (b + m) * 2 * b + b + n 53 | 54 | 55 | def list_to_flat(f_hat_list): 56 | """ 57 | A function on the SO(3) spectrum can be represented as: 58 | 1. a list f_hat of matrices f_hat[l] of size (2l+1, 2l+1) 59 | 2. a flat vector which is the concatenation of the flattened matrices 60 | 3. a zero-padded tensor with axes l, m, n. 61 | 62 | This function converts 1 to 2. 63 | 64 | :param f_hat: a list of matrices 65 | :return: a flat vector 66 | """ 67 | return np.hstack([a.flat for a in f_hat_list]) 68 | 69 | 70 | def num_spectral_coeffs_up_to_order(b): 71 | """ 72 | The SO(3) spectrum consists of matrices of size (2l+1, 2l+1) for l=0, ..., b - 1. 73 | This function computes the number of elements in a spectrum up to (but excluding) b - 1. 74 | 75 | The number of elements up to and including order L is 76 | N_L = sum_{l=0}^L (2l+1)^2 = 1/3 (2 L + 1) (2 L + 3) (L + 1) 77 | 78 | :param b: bandwidth 79 | :return: the number of spectral coefficients 80 | """ 81 | L_max = b - 1 82 | assert L_max >= 0 83 | return ((2 * L_max + 1) * (2 * L_max + 3) * (L_max + 1)) // 3 84 | 85 | 86 | def flat_to_list(f_hat_flat, b): 87 | """ 88 | A function on the SO(3) spectrum can be represented as: 89 | 1. a list f_hat of matrices f_hat[l] of size (2l+1, 2l+1) 90 | 2. a flat vector which is the concatenation of the flattened matrices 91 | 3. a zero-padded tensor with axes l, m, n. 92 | 93 | This function converts 2 to 1. 94 | 95 | :param f_hat: a flat vector 96 | :return: a list of matrices 97 | """ 98 | f_hat_list = [] 99 | start = 0 100 | for l in range(b): 101 | f_hat_list.append(f_hat_flat[start:start + (2 * l + 1) ** 2].reshape(2 * l + 1, 2 * l + 1)) 102 | start += (2 * l + 1) ** 2 103 | return f_hat_list 104 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/irrep_bases.pyx: -------------------------------------------------------------------------------- 1 | """ 2 | There are a number of different bases for the irreducible representations of SO(3), 3 | each of which results in a different form for the irrep matrices. 4 | This file contains routines that produce change-of-basis matrices 5 | to take you from one basis to the others. 6 | 7 | Recall that all irreducible representations of SO(3) appear in the 8 | decomposition of the regular representations on well-behaved functions 9 | f: S^2 -> C or f : S^2 -> R 10 | from the sphere S^2 to the real or complex numbers. 11 | 12 | The regular representation is defined by left translation: 13 | (T(g) f)(h) = f(g^{-1} h) 14 | 15 | The most common basis for the irreducible representation of weight l are some 16 | form of *complex* spherical harmonics (CSH) Y_l^m, for -l <= m <= l. 17 | 18 | For real functions, one can use real spherical harmonics (RSH) S_l^m, 19 | which have the same indexing scheme and are related to the CSH 20 | by a unitary change of basis. 21 | 22 | For both CSH and RSH, there are a number of normalization conventions, 23 | as described in spherical_harmonics.py and in [1]. However, these differ 24 | by either 25 | 1) a constant scale factor of sqrt(4 pi), or 26 | 2) a scale factor (-1)^m, which is the same for +m and -m. 27 | Since the RSH S_l^m is obtained by a linear combination of complex Y_l^m and Y_l^{-m} (see [1]), 28 | the process of changing normalization and that of changing CSH to RSH commute (we can pull out the scale/phase factor). 29 | Since the CSH-RSH change of basis is a unitary transformation, the change of basis maps each kind of CSH to a kind of 30 | RSH that has the same normalization properties. 31 | 32 | When changing the normalization, the change-of-basis matrix need not be unitary. 33 | In particular, all changes in normalization, except quantum <--> seismology, lead to non-unitary matrices. 34 | 35 | Besides normalization, the harmonics can be rearanged in different orders than m=-l,...,l 36 | This is useful because the Pinchon-Hoggan J matrix assumes a block structure in a certain ordering. 37 | 38 | For each normalization convention, we have the following bases: 39 | - Complex centered (cc): Y^{-l}, ..., Y^{l} 40 | - Real centered (rc): S^{-l}, ..., S^{l} 41 | - Real block Pinchon-Hoggan (rb): this basis is aligned with the subspaces 42 | E_xyz,k (etc.) described by Pinchon & Hoggan, and is obtained by a reordering of the RSH. 43 | In this basis, the Pinchon-Hoggan J matrix has a block structure. 44 | 45 | References: 46 | [1] http://en.wikipedia.org/wiki/Spherical_harmonics#Conventions 47 | [2] Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes. 48 | """ 49 | 50 | import numpy as np 51 | cimport numpy as np 52 | import collections 53 | from scipy.linalg import block_diag 54 | 55 | INT_TYPE = np.int64 56 | ctypedef np.int64_t INT_TYPE_t 57 | 58 | FLOAT_TYPE = np.float64 59 | ctypedef np.float64_t FLOAT_TYPE_t 60 | 61 | COMPLEX_TYPE = np.complex128 62 | ctypedef np.complex128_t COMPLEX_TYPE_t 63 | 64 | 65 | def change_of_basis_matrix(l, frm=('complex', 'seismology', 'centered', 'cs'), to=('real', 'quantum', 'centered', 'cs')): 66 | """ 67 | Compute change-of-basis matrix that takes the 'frm' basis to the 'to' basis. 68 | Each basis is identified by: 69 | 1) A field (real or complex) 70 | 2) A normalization / phase convention ('seismology', 'quantum', 'nfft', or 'geodesy') 71 | 3) An ordering convention ('centered', 'block') 72 | 4) Whether to use Condon-Shortley phase (-1)^m for m > 0 ('cs', 'nocs') 73 | 74 | Let B = change_of_basis_matrix(l, frm, to). 75 | Then if Y is a vector in the frm basis, B.dot(Y) represents the same vector in the to basis. 76 | 77 | :param l: the weight (non-negative integer) of the irreducible representation, or an iterable of weights. 78 | :param frm: a 3-tuple (field, normalization, ordering) indicating the input basis. 79 | :param to: a 3-tuple (field, normalization, ordering) indicating the output basis. 80 | :return: a (2 * l + 1, 2 * l + 1) change of basis matrix. 81 | """ 82 | from_field, from_normalization, from_ordering, from_cs = frm 83 | to_field, to_normalization, to_ordering, to_cs = to 84 | 85 | if isinstance(l, collections.Iterable): 86 | blocks = [change_of_basis_matrix(li, frm, to) 87 | for li in l] 88 | return block_diag(*blocks) 89 | 90 | # First, bring us to the centered basis: 91 | if from_ordering == 'block': 92 | B = _c2b(l).T 93 | elif from_ordering == 'centered': 94 | B = np.eye(2 * l + 1) 95 | else: 96 | raise ValueError('Invalid from_order: ' + str(from_ordering)) 97 | 98 | # Make sure we're using CS-phase (this should work for both real and complex bases) 99 | if from_cs == 'nocs': 100 | m = np.arange(-l, l + 1) 101 | B = ((-1.) ** (m * (m > 0)))[:, None] * B 102 | elif from_cs != 'cs': 103 | raise ValueError('Invalid from_cs: ' + str(from_cs)) 104 | 105 | # If needed, change complex to real or real to complex 106 | # (we know how to do that in the centered, CS-phase bases) 107 | if from_field != to_field: 108 | if from_field == 'complex' and to_field == 'real': 109 | B = _cc2rc(l).dot(B) 110 | elif from_field == 'real' and to_field == 'complex': 111 | B = _cc2rc(l).conj().T.dot(B) 112 | else: 113 | raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field)) 114 | 115 | # If needed, change the normalization: 116 | if from_normalization != to_normalization: 117 | # First, change normalization to quantum 118 | if from_normalization == 'seismology': 119 | B = _seismology2quantum(l, full_matrix=False)[:, None] * B 120 | elif from_normalization == 'geodesy': 121 | B = _geodesy2quantum(l, full_matrix=False)[:, None] * B 122 | elif from_normalization == 'nfft': 123 | B = _nfft2quantum(l, full_matrix=False)[:, None] * B 124 | elif from_normalization != 'quantum': 125 | raise ValueError('Invalud from_normalization:' + str(from_normalization)) 126 | 127 | # We're now in quantum normalization, change to output normalization 128 | if to_normalization == 'seismology': 129 | B = (1. / _seismology2quantum(l, full_matrix=False))[:, None] * B 130 | elif to_normalization == 'geodesy': 131 | B = (1. / _geodesy2quantum(l, full_matrix=False))[:, None] * B 132 | elif to_normalization == 'nfft': 133 | B = (1. / _nfft2quantum(l, full_matrix=False))[:, None] * B 134 | elif to_normalization != 'quantum': 135 | raise ValueError('Invalid to_normalization:' + str(to_normalization)) 136 | 137 | #if from_field != to_field: 138 | # if from_field == 'complex' and to_field == 'real': 139 | # B = cc2rc(l).dot(B) 140 | # elif from_field == 'real' and to_field == 'complex': 141 | # #B = cc2rc(l).conj().T.dot(B) 142 | # pass 143 | # else: 144 | # raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field)) 145 | #if to_field == 'real': 146 | # B = cc2rc(l).dot(B) 147 | #elif to_field != 'complex': 148 | # raise ValueError('Invalid to_field: ' + str(to_field)) 149 | 150 | # Set the correct CS phase 151 | if to_cs == 'nocs': 152 | # We're in CS phase now, so cancel it: 153 | m = np.arange(-l, l + 1) 154 | B = ((-1.) ** (m * (m > 0)))[:, None] * B 155 | elif to_cs != 'cs': 156 | raise ValueError('Invalid to_cs: ' + str(to_cs)) 157 | 158 | # If needed, change the order from centered: 159 | if to_ordering == 'block': 160 | B = _c2b(l).dot(B) 161 | elif to_ordering != 'centered': 162 | raise ValueError('Invalid to_ordering:' + str(to_ordering)) 163 | 164 | return B 165 | 166 | 167 | #TODO: make sure that change_of_basis_function accepts matrices, where each row is a vector to be changed of basis. 168 | def change_of_basis_function(l, frm=('complex', 'seismology', 'centered', 'cs'), 169 | to=('real', 'quantum', 'centered', 'cs')): 170 | """ 171 | Return a function that will compute the change-of-basis that takes the 'frm' basis to the 'to' basis. 172 | Each basis is identified by: 173 | 1) A field (real or complex) 174 | 2) A normalization / phase convention ('seismology', 'quantum', or 'geodesy') 175 | 3) An ordering convention ('centered', 'block') 176 | 4) Whether to use Condon-Shortley phase (-1)^m for m > 0 ('cs', 'nocs') 177 | 178 | :param l: the weight (non-negative integer) of the irreducible representation, or an iterable of weights. 179 | :param frm: a 3-tuple (field, normalization, ordering) indicating the input basis. 180 | :param to: a 3-tuple (field, normalization, ordering) indicating the output basis. 181 | :return: 182 | """ 183 | 184 | from_field, from_normalization, from_ordering, from_cs = frm 185 | to_field, to_normalization, to_ordering, to_cs = to 186 | 187 | if not isinstance(l, np.ndarray): # collections.Iterable): 188 | l = np.atleast_1d(np.array(l)) 189 | 190 | # First, bring us to the centered basis: 191 | if from_ordering == 'block': 192 | f1 = _b2c_func(l) 193 | elif from_ordering == 'centered': 194 | f1 = lambda x: x 195 | else: 196 | raise ValueError('Invalid from_order: ' + str(from_ordering)) 197 | 198 | ms = np.zeros(np.sum(2 * l + 1), dtype=INT_TYPE) 199 | ls = np.zeros(np.sum(2 * l + 1), dtype=INT_TYPE) 200 | i = 0 201 | for ll in l: 202 | for mm in range(-ll, ll + 1): 203 | ms[i] = mm 204 | ls[i] = ll 205 | i += 1 206 | 207 | # Make sure we're using CS-phase (this should work for both real and complex bases) 208 | if from_cs == 'nocs': 209 | p = ((-1.) ** (ms * (ms > 0))) 210 | f2 = lambda x: f1(x) * p 211 | elif from_cs == 'cs': 212 | f2 = f1 213 | else: # elif from_cs != 'cs': 214 | raise ValueError('Invalid from_cs: ' + str(from_cs)) 215 | 216 | # If needed, change complex to real or real to complex 217 | # (we know how to do that in the centered, CS-phase bases) 218 | if from_field != to_field: 219 | if from_field == 'complex' and to_field == 'real': 220 | #B = _cc2rc(l).dot(B) 221 | #pos_m = m > 0 222 | #neg_m = m < 0 223 | #zero_m = m == 0 224 | #f3 = lambda x: r(f2(x) * _cc2rc_func(x, m), 3) 225 | f3 = lambda x: _cc2rc_func(f2(x), ms) 226 | elif from_field == 'real' and to_field == 'complex': 227 | f3 = lambda x: _rc2cc_func(f2(x), ms) 228 | #raise NotImplementedError('Real to complex not implemented yet') 229 | else: 230 | raise ValueError('Invalid field:' + str(from_field) + ', ' + str(to_field)) 231 | else: 232 | f3 = f2 233 | 234 | # If needed, change the normalization: 235 | if from_normalization != to_normalization: 236 | # First, change normalization to quantum 237 | if from_normalization == 'seismology': 238 | f4 = lambda x: f3(x) * _seismology2quantum(l, full_matrix=False) 239 | elif from_normalization == 'geodesy': 240 | f4 = lambda x: f3(x) * _geodesy2quantum(l, full_matrix=False) 241 | elif from_normalization == 'nfft': 242 | f4 = lambda x: f3(x) * _nfft2quantum(l, full_matrix=False) 243 | elif from_normalization == 'quantum': 244 | f4 = f3 245 | else: # elif from_normalization != 'quantum': 246 | raise ValueError('Invalud from_normalization:' + str(from_normalization)) 247 | 248 | # We're now in quantum normalization, change to output normalization 249 | if to_normalization == 'seismology': 250 | f5 = lambda x: f4(x) / _seismology2quantum(l, full_matrix=False) 251 | elif to_normalization == 'geodesy': 252 | f5 = lambda x: f4(x) / _geodesy2quantum(l, full_matrix=False) 253 | elif to_normalization == 'nfft': 254 | f5 = lambda x: f4(x) / _nfft2quantum(l, full_matrix=False) 255 | elif to_normalization == 'quantum': 256 | f5 = f4 257 | else: # elif to_normalization != 'quantum': 258 | raise ValueError('Invalid to_normalization:' + str(to_normalization)) 259 | else: 260 | f5 = f3 261 | 262 | # Set the correct CS phase 263 | if to_cs == 'nocs': 264 | # We're in CS phase now, so cancel it: 265 | #m = np.arange(-l, l + 1) 266 | #B = ((-1.) ** (m * (m > 0)))[:, None] * B 267 | p = ((-1.) ** (ms * (ms > 0))) 268 | f6 = lambda x: f5(x) * p 269 | elif to_cs == 'cs': 270 | f6 = f5 271 | elif to_cs != 'cs': 272 | raise ValueError('Invalid to_cs: ' + str(to_cs)) 273 | 274 | # If needed, change the order from centered: 275 | if to_ordering == 'block': 276 | #B = _c2b(l).dot(B) 277 | #raise NotImplementedError('Block basis not supported yet') 278 | f7 = lambda x: _c2b_func(l)(f6(x)) 279 | elif to_ordering == 'centered': 280 | f7 = f6 281 | else: 282 | raise ValueError('Invalid to_ordering:' + str(to_ordering)) 283 | 284 | return f7 285 | 286 | 287 | def _cc2rc(l): 288 | """ 289 | Compute change of basis matrix from the complex centered (cc) basis 290 | to the real centered (rc) basis. 291 | 292 | Let Y be a vector of complex spherical harmonics: 293 | Y = (Y^{-l}, ..., Y^0, ..., Y^l)^T 294 | Let S be a vector of real spherical harmonics as defined on the SH wiki page: 295 | S = (S^{-l}, ..., S^0, ..., S^l)^T 296 | Let B = cc2rc(l) 297 | Then S = B.dot(Y) 298 | 299 | B is a complex unitary matrix. 300 | 301 | Formula taken from: 302 | http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2 303 | """ 304 | 305 | B = np.zeros((2 * l + 1, 2 * l + 1), dtype=complex) 306 | for m in range(-l, l + 1): 307 | for n in range(-l, l + 1): 308 | row_ind = m + l 309 | col_ind = n + l 310 | if m == 0 and n == 0: 311 | B[row_ind, col_ind] = np.sqrt(2) 312 | if m > 0 and m == n: 313 | B[row_ind, col_ind] = (-1.) ** m 314 | elif m > 0 and m == -n: 315 | B[row_ind, col_ind] = 1. 316 | elif m < 0 and m == n: 317 | B[row_ind, col_ind] = 1j 318 | elif m < 0 and m == -n: 319 | B[row_ind, col_ind] = -1j * ((-1.) ** m) 320 | 321 | return (1.0 / np.sqrt(2)) * B 322 | 323 | 324 | def _cc2rc_func(np.ndarray[COMPLEX_TYPE_t, ndim=1] x, 325 | np.ndarray[INT_TYPE_t, ndim=1] m_arr): 326 | """ 327 | Compute change of basis from the complex centered (cc) basis 328 | to the real centered (rc) basis. 329 | 330 | Let Y be a vector of complex spherical harmonics: 331 | Y = (Y^{-l}, ..., Y^0, ..., Y^l)^T 332 | Let S be a vector of real spherical harmonics as defined on the SH wiki page: 333 | S = (S^{-l}, ..., S^0, ..., S^l)^T 334 | Let B = cc2rc(l) 335 | Then S = B.dot(Y) 336 | 337 | B is a complex unitary matrix. 338 | 339 | Formula taken from: 340 | http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2 341 | """ 342 | 343 | cdef int i = 0 344 | cdef np.ndarray[FLOAT_TYPE_t, ndim=1] x_out = np.empty(x.size) 345 | cdef double sq2 = np.sqrt(2) 346 | cdef double isq2 = 1. / sq2 347 | 348 | for i in range(m_arr.size): 349 | m = m_arr[i] 350 | if m > 0: 351 | x_out[i] = ((-1.) ** m * x[i] + x[i - 2 * m]).real * isq2 352 | elif m < 0: 353 | x_out[i] = (1j * x[i] - 1j * ((-1.) ** m) * x[i - 2 * m]).real * isq2 354 | else: 355 | x_out[i] = x[i].real 356 | 357 | return x_out 358 | 359 | 360 | def _rc2cc_func(np.ndarray[FLOAT_TYPE_t, ndim=1] x, 361 | np.ndarray[INT_TYPE_t, ndim=1] m_arr): 362 | """ 363 | Compute change of basis from the real centered (rc) basis 364 | to the complex centered (cc) basis. 365 | 366 | Formula taken from: 367 | http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form_2 368 | """ 369 | 370 | cdef int i = 0 371 | cdef np.ndarray[COMPLEX_TYPE_t, ndim=1] x_out = np.empty(x.size, dtype=COMPLEX_TYPE) 372 | cdef double sq2 = np.sqrt(2) 373 | cdef double isq2 = 1. / sq2 374 | 375 | for i in range(m_arr.size): 376 | m = m_arr[i] 377 | if m > 0: 378 | x_out[i] = ((-1.) ** m * x[i - 2 * m] * 1j + (-1.) ** m * x[i]) * isq2 379 | elif m < 0: 380 | x_out[i] = (-1j * x[i] + x[i - 2 * m]) * isq2 381 | else: 382 | x_out[i] = x[i] 383 | 384 | return x_out 385 | 386 | 387 | def _c2b(l, full_matrix=True): 388 | """ 389 | Compute change of basis matrix from the centered basis to 390 | the Pinchon-Hoggan block basis, in which the Pinchon-Hoggan J matrices 391 | are brought in block form. 392 | 393 | Let B = c2b(l) 394 | then B.dot(J_l).dot(B.T) is in block form with 4 blocks, 395 | as described by PH. 396 | """ 397 | k = int(l) // 2 398 | if l % 2 == 0: 399 | # Permutation as defined by Pinchon-Hoggan for 1-based indices, 400 | # and l = 2 k 401 | sigma = np.array([2 * i for i in range(1, 2 * k + 1)] 402 | + [2 * i - 1 for i in range(1, 2 * k + 2)]) 403 | else: 404 | # Permutation as defined by Pinchon-Hoggan for 1-based indices, 405 | # and l = 2 k + 1 406 | sigma = np.array([2 * i for i in range(1, 2 * k + 2)] 407 | + [2 * i - 1 for i in range(1, 2 * k + 3)]) 408 | 409 | if full_matrix: 410 | # From permutation array sigma, create a permutation matrix B: 411 | B = np.zeros((2 * l + 1, 2 * l + 1)) 412 | B[np.arange(2 * l + 1), sigma - 1] = 1. 413 | return B 414 | else: 415 | return sigma 416 | 417 | 418 | def _c2b_func(l): 419 | """ 420 | 421 | :param l: 422 | :return: 423 | """ 424 | sigma = np.hstack([_c2b(li, full_matrix=False) - 1 for li in l]) 425 | i_begin = 0 426 | for li in l: 427 | sigma[i_begin:i_begin + 2 * li + 1] += i_begin 428 | i_begin += 2 * li + 1 429 | f = lambda x: x[sigma] 430 | return f 431 | 432 | 433 | def _b2c_func(l): 434 | sigma = np.hstack([_c2b(li, full_matrix=False) - 1 for li in l]) 435 | i_begin = 0 436 | for li in l: 437 | sigma[i_begin:i_begin + 2 * li + 1] += i_begin 438 | i_begin += 2 * li + 1 439 | sigma_inv = np.argsort(sigma) 440 | f = lambda x: x[sigma_inv] 441 | return f 442 | 443 | 444 | 445 | 446 | def _seismology2quantum(l, full_matrix=False): 447 | """ 448 | 449 | :param l: 450 | :param full_matrix: 451 | :return: 452 | """ 453 | if isinstance(l, collections.Iterable): 454 | diags = [_seismology2quantum(li, full_matrix=False) for li in l] 455 | diagonal = np.hstack(diags) 456 | 457 | if full_matrix: 458 | return np.diag(diagonal) 459 | else: 460 | return diagonal 461 | 462 | diagonal = (-np.ones(2 * l + 1)) ** np.arange(-l, l + 1) 463 | if full_matrix: 464 | return np.diag(diagonal) 465 | else: 466 | return diagonal 467 | 468 | 469 | def _geodesy2quantum(l, full_matrix=False): 470 | if isinstance(l, collections.Iterable): 471 | diags = [_geodesy2quantum(li, full_matrix=False) for li in l] 472 | diagonal = np.hstack(diags) 473 | 474 | if full_matrix: 475 | return np.diag(diagonal) 476 | else: 477 | return diagonal 478 | 479 | diagonal = (-np.ones(2 * l + 1)) ** np.arange(-l, l + 1) 480 | diagonal /= np.sqrt(4 * np.pi) 481 | if full_matrix: 482 | return np.diag(diagonal) 483 | else: 484 | return diagonal 485 | 486 | 487 | def _nfft2quantum(l, full_matrix=False): 488 | 489 | if isinstance(l, collections.Iterable): 490 | diags = [_nfft2quantum(li, full_matrix=False) for li in l] 491 | diagonal = np.hstack(diags) 492 | 493 | if full_matrix: 494 | return np.diag(diagonal) 495 | else: 496 | return diagonal 497 | 498 | diagonal = np.ones(2 * l + 1) * np.sqrt((2 * l + 1) / (4. * np.pi)) 499 | # nfft only has (-1)^m phase factor for positive m; quantum has both pos and neg, so add phase factor to neg inds: 500 | # -> this is now done using a CS setting 501 | #m = np.arange(-l, l + 1) 502 | #diagonal *= ((-1) ** (m * (m < 0))) 503 | 504 | m = np.arange(-l, l + 1) 505 | diagonal *= (-1.) ** m 506 | 507 | if full_matrix: 508 | return np.diag(diagonal) 509 | else: 510 | return diagonal 511 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/J_block_0-150.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/representations/SO3/pinchon_hoggan/J_block_0-150.npy -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/J_dense_0-150.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/representations/SO3/pinchon_hoggan/J_dense_0-150.npy -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/representations/SO3/pinchon_hoggan/__init__.py -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/download.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | import requests 4 | 5 | 6 | def download(url): 7 | base = os.path.basename(url) 8 | path = os.path.join(os.path.dirname(__file__), base) 9 | if not os.path.isfile(path): 10 | with open(path, 'wb') as f: 11 | f.write(requests.get(url).content) 12 | return np.load(path, encoding='latin1', allow_pickle=True) 13 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan.pyx: -------------------------------------------------------------------------------- 1 | """ 2 | Code for rotating spherical harmonics expansions by the method described in: 3 | Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes 4 | D. Pinchon, P. E. Hoggan 5 | 6 | All functions in this file assume that we're working in the basis of real, quantum-normalized, Pinchon-Hoggan block 7 | spherical harmonics 8 | 9 | This is NOT a user-facing API, so the interface of this file may change. 10 | """ 11 | 12 | import numpy as np 13 | cimport numpy as np 14 | cimport cython 15 | 16 | from lie_learn.broadcasting import generalized_broadcast 17 | 18 | # Load the J-matrices, which are stored in the same folder as this file 19 | #import os 20 | #Jb = np.load(os.path.join(os.path.dirname(__file__), 'J_block_0-478.npy'), allow_pickle=True) 21 | 22 | FLOAT_TYPE = np.float64 23 | ctypedef np.float64_t FLOAT_TYPE_t 24 | INT_TYPE = np.int 25 | ctypedef np.int64_t INT_TYPE_t 26 | 27 | 28 | def apply_rotation_block(g, X, irreps, c2b, J_block, l_max, X_out=None): 29 | 30 | X, g = generalized_broadcast([X, g]) 31 | out_shape = X.shape 32 | X = X.reshape(-1, X.shape[-1]).copy() 33 | g = g.reshape(-1, g.shape[-1]).copy() 34 | 35 | if X_out is None: 36 | X_out = np.empty_like(X) 37 | X_out = X_out.reshape(-1, X.shape[-1]) 38 | X_temp = np.empty_like(X_out) 39 | 40 | apply_z_rotation_block(g[:, 2], X, irreps, c2b, l_max, X_out=X_out) 41 | apply_J_block(X_out, J_block, X_out=X_temp) 42 | apply_z_rotation_block(g[:, 1], X_temp, irreps, c2b, l_max, X_out=X_out) 43 | apply_J_block(X_out, J_block, X_out=X_temp) 44 | apply_z_rotation_block(g[:, 0], X_temp, irreps, c2b, l_max, X_out=X_out) 45 | 46 | return X_out.reshape(out_shape) 47 | 48 | 49 | @cython.wraparound(False) 50 | @cython.nonecheck(False) 51 | @cython.boundscheck(False) 52 | cdef apply_z_rotation_block(np.ndarray[FLOAT_TYPE_t, ndim=1] angles, 53 | np.ndarray[FLOAT_TYPE_t, ndim=2] X, 54 | np.ndarray[INT_TYPE_t, ndim=1] irreps, 55 | np.ndarray[INT_TYPE_t, ndim=1] c2b, 56 | int l_max, 57 | np.ndarray[FLOAT_TYPE_t, ndim=2] X_out): 58 | """ 59 | Apply the rotation about the z-axis by angle angles[i] to the vector 60 | X[i, :] for all i. The vectors in X are assumed to be in the basis of real 61 | block spherical harmonics, corresponding to the irreps. 62 | 63 | In the *centered* basis, the z-axis rotation matrix (called X(angle) in the P&H paper) has a special 64 | form with cosines of different frequencies on the diagonal and sines on the 65 | anti-diagonal. This matrix is very sparse, so constructing it explicitly is 66 | very inefficient. This function applies a batch of such 'cross' matrices, 67 | represented implicitly by the corresponding angles, to a batch of vectors, 68 | without explicitly constructing the matrices. 69 | To do this in the block basis, we use a permutation array c2b that takes an index in the centered basis 70 | and returns the index for the block basis (which is a permutation of the centered basis). 71 | 72 | Args: 73 | angles: matrix of angles, shape (num_data, num_angles) 74 | irreps: a list of irrep weights (integers) 75 | c2b: an array of length dim, where c2b[i] is the index of centere basis vector i in the block basis 76 | l_max: an integer that is equal to np.max(irreps) 77 | X: matrix to be rotated, shape (num_data, dim) 78 | X_out: matrix where output will be stored, shape (dim, num_data, num_angles) 79 | 80 | Returns: 81 | X_out: matrix of shape (dim, num_data, num_angles), such that the 82 | vector X_out[:, i, j] is the rotation of X[:,i] by angles[i, j]. 83 | """ 84 | 85 | cdef int irrep_ind # Index into the IRREPS2 matrix 86 | cdef int irrep # The irrep weight 87 | cdef int center # The index of the center element (l=0) of current irrep 88 | cdef int offset # The offset of the current coordinate from the center 89 | cdef int abs_offset # The absolute value of the offset 90 | cdef int offset_sign # The sign of the offset 91 | #cdef int l_max = np.max(irreps) 92 | 93 | cdef int Xs0 = X.shape[0] 94 | cdef int Xs1 = X.shape[1] 95 | #cdef int Cs2 96 | 97 | cdef int vec 98 | cdef int coord 99 | cdef int angle_ind 100 | 101 | cdef int coord1_block 102 | cdef int coord2_block 103 | 104 | cdef np.ndarray[FLOAT_TYPE_t, ndim=2] C = np.cos(angles[None, :] * np.arange(l_max + 1)[:, None]) 105 | cdef np.ndarray[FLOAT_TYPE_t, ndim=2] S = np.sin(angles[None, :] * np.arange(l_max + 1)[:, None]) 106 | #Cs2 = C.shape[2] 107 | 108 | # Check that the irrep dimensions sum to the right dimensionality 109 | #assert (2 * irreps + 1).sum() == X.shape[0], "Invalid irrep dimensions" 110 | #assert angles.shape[0] == X.shape[1] 111 | #assert X_out.shape[0] == X.shape[0], "Invalid shape for X_out" 112 | #assert X_out.shape[1] == X.shape[1], "Invalid shape for X_out" 113 | #assert X_out.shape[2] == angles.shape[1], "Invalid shape for X_out" 114 | 115 | for vec in range(Xs0): # Xs1): 116 | irrep_ind = 0 # Start at the first irrep 117 | irrep = irreps[irrep_ind] 118 | center = irrep 119 | for coord in range(Xs1): # Xs0): # X.shape[0]): 120 | offset = coord - center 121 | if offset > irrep: # Finished with the current irrep? 122 | irrep_ind += 1 # Go to the next irrep 123 | irrep = irreps[irrep_ind] # Get its weight from the list 124 | center = coord + irrep # Compute the new center 125 | offset = -irrep # equivalent to offset=coord-center; 126 | 127 | # Compute the absolute value and sign of the offset 128 | abs_offset = abs(offset) 129 | if offset >= 0: 130 | offset_sign = 1 131 | else: 132 | offset_sign = -1 133 | 134 | coord1_block = c2b[coord] 135 | coord2_block = c2b[center - offset] 136 | 137 | # Compute the value of the transformed coordinate 138 | # Note: we're always adding *two* values, even when offset=0 and hence there is only 139 | # one non-zero element in that row. This is not a problem because S2(0, vec) is always 0. 140 | #X_out[coord1_block, vec] = C[abs_offset, vec] * X[coord1_block, vec] \ 141 | # - offset_sign * S[abs_offset, vec] * X[coord2_block, vec] 142 | X_out[vec, coord1_block] = C[abs_offset, vec] * X[vec, coord1_block] \ 143 | - offset_sign * S[abs_offset, vec] * X[vec, coord2_block] 144 | 145 | 146 | return X_out 147 | 148 | 149 | @cython.wraparound(False) 150 | @cython.nonecheck(False) 151 | @cython.boundscheck(False) 152 | cdef apply_J_block(np.ndarray[FLOAT_TYPE_t, ndim=2] X, 153 | list J_block, 154 | np.ndarray[FLOAT_TYPE_t, ndim=2] X_out): 155 | """ 156 | Multiply the Pinchon-Hoggan J matrix by a matrix X. 157 | 158 | This function uses the J matrix in the Pinchon-Hoggan block-basis. In this basis, the J-matrix is a block matrix 159 | with 4 blocks. This function performs this block-multiplication efficiently. 160 | 161 | :param X: numpy array of shape (dim, N) where dim is the total dimension of the representation and N is the number 162 | of vectors to be multiplied by J. 163 | :param J_block: list of list of precomputed J matrices in block form. 164 | :return: 165 | """ 166 | 167 | cdef int l 168 | cdef int k 169 | cdef int rep_begin = 0 170 | cdef int b1s 171 | cdef int b1e 172 | cdef int b2s 173 | cdef int b2e 174 | cdef int b3s 175 | cdef int b3e 176 | cdef int b4s 177 | cdef int b4e 178 | cdef int li 179 | 180 | # Loop over irreps 181 | for li in range(len(J_block)): 182 | 183 | Jl = J_block[li] 184 | k = Jl[0].shape[0] 185 | l = (Jl[0].shape[0] + 2 * Jl[1].shape[0] + Jl[2].shape[0]) // 2 186 | 187 | # Determine block begin and end indices 188 | if l % 2 == 0: 189 | # blocks have dimension k, k, k, k+1 190 | b1s = rep_begin 191 | b1e = rep_begin + k 192 | b2s = rep_begin + k 193 | b2e = rep_begin + 2 * k 194 | b3s = rep_begin + 2 * k 195 | b3e = rep_begin + 3 * k 196 | b4s = rep_begin + 3 * k 197 | b4e = rep_begin + 4 * k + 1 198 | 199 | #b1 = Jbl[0:k, 0:k] 200 | #b2 = Jbl[k:2 * k, 2 * k:3 * k] 201 | ##b3 = Jbl[2 * k: 3 * k, k:2 * k] 202 | #b4 = Jbl[3 * k:, 3 * k:] 203 | else: 204 | # blocks have dimension k, k+1, k+1, k+1 205 | b1s = rep_begin 206 | b1e = rep_begin + k 207 | b2s = rep_begin + k 208 | b2e = rep_begin + 2 * k + 1 209 | b3s = rep_begin + 2 * k + 1 210 | b3e = rep_begin + 3 * k + 2 211 | b4s = rep_begin + 3 * k + 2 212 | b4e = rep_begin + 4 * k + 3 213 | 214 | #b1 = Jbl[0:k, 0:k] 215 | #b2 = Jbl[k:2 * k + 1, 2 * k + 1:3 * k + 2] 216 | ##b3 = Jbl[2 * k + 1:3 * k + 2, k:2 * k + 1] 217 | #b4 = Jbl[3 * k + 2:, 3 * k + 2:] 218 | 219 | # Multiply each block: 220 | X_out[:, b1s:b1e] = np.dot(Jl[0], X[:, b1s:b1e].T).T 221 | X_out[:, b2s:b2e] = np.dot(Jl[1], X[:, b3s:b3e].T).T 222 | X_out[:, b3s:b3e] = np.dot(Jl[1].T, X[:, b2s:b2e].T).T 223 | X_out[:, b4s:b4e] = np.dot(Jl[2], X[:, b4s:b4e].T).T 224 | 225 | rep_begin += 2 * l + 1 226 | 227 | return X_out 228 | 229 | 230 | def make_c2b(irreps): 231 | # Centered to block basis 232 | 233 | # Maybe put this in separate file irrep_bases? 234 | c2b = np.empty((2 * irreps + 1).sum(), dtype=INT_TYPE) 235 | irrep_begin = 0 236 | for l in irreps: 237 | 238 | k = int(l) / 2 239 | if l % 2 == 0: 240 | # Permutation as defined by Pinchon-Hoggan for 1-based indices, 241 | # and l = 2 k 242 | sigma = np.array([2 * i for i in range(1, 2 * k + 1)] 243 | + [2 * i - 1 for i in range(1, 2 * k + 2)]) 244 | else: 245 | # Permutation as defined by Pinchon-Hoggan for 1-based indices, 246 | # and l = 2 k + 1 247 | sigma = np.array([2 * i for i in range(1, 2 * k + 2)] 248 | + [2 * i - 1 for i in range(1, 2 * k + 3)]) 249 | 250 | sigma_inv = np.arange(0, 2 * l + 1)[np.argsort(sigma)] 251 | 252 | c2b[irrep_begin:irrep_begin + 2 * l + 1] = sigma_inv + irrep_begin 253 | irrep_begin += 2 * l + 1 254 | return c2b -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan_dense.py: -------------------------------------------------------------------------------- 1 | import os 2 | import numpy as np 3 | from scipy.linalg import block_diag 4 | 5 | # This code is not very optimized, 6 | # and can never become very efficient because it cannot exploit the sparsity of the J matrix. 7 | 8 | # Load the J-matrices, which are stored in the same folder as this file 9 | from .download import download 10 | 11 | # J matrices come from this paper 12 | # Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes 13 | # Didier Pinchon1 and Philip E Hoggan2 14 | # https://iopscience.iop.org/article/10.1088/1751-8113/40/7/011/ 15 | 16 | # Jd = download('https://github.com/AMLab-Amsterdam/lie_learn/releases/download/v1.0/J_dense_0-278.npy') 17 | base = 'J_dense_0-150.npy' 18 | path = os.path.join(os.path.dirname(__file__), base) 19 | Jd = np.load(path, allow_pickle=True) 20 | 21 | def SO3_irreps(g, irreps): 22 | global Jd 23 | 24 | # First, compute sinusoids at all required frequencies, i.e. 25 | # cos(n x) for n=0, ..., max(irreps) 26 | # sin(n x) for n=-max(irreps), ..., max(irreps) 27 | # where x ranges over the three parameters of SO(3). 28 | 29 | # In theory, it may be faster to evaluate cos(x) once and then use 30 | # Chebyshev polynomials to obtain cos(n*x), but in practice this appears 31 | # to be slower than just evaluating cos(n*x). 32 | dim = np.sum(2 * np.array(irreps) + 1) 33 | T = np.empty((dim, dim, g.shape[1])) 34 | for i in range(g.shape[1]): 35 | T[:, :, i] = block_diag(*[rot_mat(g[0, i], g[1, i], g[2, i], l, Jd[l]) for l in irreps]) 36 | return T 37 | 38 | 39 | def SO3_irrep(g, l): 40 | global Jd 41 | g = np.atleast_2d(g) 42 | T = np.empty((2 * l + 1, 2 * l + 1, g.shape[1])) 43 | for i in range(g.shape[1]): 44 | T[:, :, i] = rot_mat(g[0, i], g[1, i], g[2, i], l, Jd[l]) 45 | return T # np.squeeze(T) 46 | 47 | 48 | def z_rot_mat(angle, l): 49 | """ 50 | Create the matrix representation of a z-axis rotation by the given angle, 51 | in the irrep l of dimension 2 * l + 1, in the basis of real centered 52 | spherical harmonics (RC basis in rep_bases.py). 53 | 54 | Note: this function is easy to use, but inefficient: only the entries 55 | on the diagonal and anti-diagonal are non-zero, so explicitly constructing 56 | this matrix is unnecessary. 57 | """ 58 | M = np.zeros((2 * l + 1, 2 * l + 1)) 59 | inds = np.arange(0, 2 * l + 1, 1) 60 | reversed_inds = np.arange(2 * l, -1, -1) 61 | frequencies = np.arange(l, -l - 1, -1) 62 | M[inds, reversed_inds] = np.sin(frequencies * angle) 63 | M[inds, inds] = np.cos(frequencies * angle) 64 | return M 65 | 66 | 67 | def rot_mat(alpha, beta, gamma, l, J): 68 | """ 69 | Compute the representation matrix of a rotation by ZYZ-Euler 70 | angles (alpha, beta, gamma) in representation l in the basis 71 | of real spherical harmonics. 72 | 73 | The result is the same as the wignerD_mat function by Johann Goetz, 74 | when the sign of alpha and gamma is flipped. 75 | 76 | The forementioned function is here: 77 | https://sites.google.com/site/theodoregoetz/notes/wignerdfunction 78 | """ 79 | Xa = z_rot_mat(alpha, l) 80 | Xb = z_rot_mat(beta, l) 81 | Xc = z_rot_mat(gamma, l) 82 | return Xa.dot(J).dot(Xb).dot(J).dot(Xc) 83 | 84 | 85 | def derivative_z_rot_mat(angle, l): 86 | M = np.zeros((2 * l + 1, 2 * l + 1)) 87 | inds = np.arange(0, 2 * l + 1, 1) 88 | reversed_inds = np.arange(2 * l, -1, -1) 89 | frequencies = np.arange(l, -l - 1, -1) 90 | M[inds, reversed_inds] = np.cos(frequencies * angle) * frequencies 91 | M[inds, inds] = -np.sin(frequencies * angle) * frequencies 92 | return M 93 | 94 | 95 | def derivative_rot_mat(alpha, beta, gamma, l, J): 96 | Xa = z_rot_mat(alpha, l) 97 | Xb = z_rot_mat(beta, l) 98 | Xc = z_rot_mat(gamma, l) 99 | dXa_da = derivative_z_rot_mat(alpha, l) 100 | dXb_db = derivative_z_rot_mat(beta, l) 101 | dXc_dc = derivative_z_rot_mat(gamma, l) 102 | 103 | dDda = dXa_da.dot(J).dot(Xb).dot(J).dot(Xc) 104 | dDdb = Xa.dot(J).dot(dXb_db).dot(J).dot(Xc) 105 | dDdc = Xa.dot(J).dot(Xb).dot(J).dot(dXc_dc) 106 | return dDda, dDdb, dDdc 107 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/pinchon_hoggan/pinchon_hoggan_parsing.py: -------------------------------------------------------------------------------- 1 | 2 | import os 3 | import numpy as np 4 | import scipy.sparse as sp 5 | from ..irrep_bases import change_of_basis_function, change_of_basis_matrix 6 | 7 | 8 | def convert_all(num_mat_J_folder='/Users/user/Projects/LieLearn/SO3/shrot/legacy/pinchon2/Maple/NumMatJ/', 9 | out_folder='./', l_min=0, l_max=None): 10 | """ 11 | Convert all numMatJ-XXX.dat files in a given folder to produce the following python pickle files: 12 | 13 | 1) J_dense_0-[l_max].npy: 14 | 2) J_sparse_0-[l_max].npy: 15 | 3) J_block_0-[l_max].npy 16 | 17 | The numMatJ-XXX.dat files are the output of mkNumMatJ.mw 18 | 19 | :param num_mat_J_folder: 20 | :return: 21 | """ 22 | J_dense = [] 23 | J_sparse = [] 24 | J_block = [] 25 | 26 | if l_max is None: 27 | files = [f for f in os.listdir(num_mat_J_folder) if '.dat' in f] 28 | nums = [int(f[len("numMatJ-"):-len(".dat")]) for f in files] 29 | l_max = np.max(nums) 30 | print('Maximum l found:', l_max) 31 | 32 | for l in range(l_min, l_max + 1): 33 | print('Parsing l', l) 34 | 35 | filename = os.path.join(num_mat_J_folder, 'numMatJ-' + str(l) + '.dat') 36 | 37 | # Obtain the J matrix as a dense numpy array 38 | Jd = parse_J_single_l(filename) 39 | J_dense.append(Jd) 40 | J_sparse.append(sp.csr_matrix(Jd)) 41 | 42 | print('Saving dense matrices...') 43 | np.save(os.path.join(out_folder, 'J_dense_' + str(l_min) + '-' + str(l_max)), J_dense) 44 | 45 | print('Saving sparse matrices...') 46 | np.save(os.path.join(out_folder, 'J_sparse_' + str(l_min) + '-' + str(l_max)), J_sparse) 47 | del J_sparse 48 | 49 | print('Converting to block basis...') 50 | J_block = make_block_J(J_dense) 51 | 52 | print('Saving block matrices...') 53 | np.save(os.path.join(out_folder, 'J_block_' + str(l_min) + '-' + str(l_max)), J_block) 54 | 55 | 56 | def parse_J(file): 57 | """ 58 | Loads the numMatJ.dat files provided by Pinchon & Hoggan into numpy matrices. 59 | I saved the result as a compressed pickle, so this function will 60 | only be needed when J matrices for larger l are needed. 61 | """ 62 | 63 | f = open(file) 64 | lmax = int(f.readline().split(' ')[1]) 65 | 66 | matj = [np.zeros((2 * l + 1, 2 * l + 1)) for l in range(lmax)] 67 | 68 | # Fill in matrix l=0 and l=1 69 | matj[0][0, 0] = 1.0 70 | matj[1][0, 1] = -1.0 71 | matj[1][1, 0] = -1.0 72 | matj[1][2, 2] = 1.0 73 | 74 | # Read out matrices >= 2 75 | for l in range(2, lmax): 76 | 77 | # Read and discard l-value: 78 | lval = int(f.readline().split(' ')[1]) 79 | assert lval == l 80 | 81 | k = l/2 82 | k1 = k 83 | for i in range(k): 84 | for j in range(i): 85 | x = float(f.readline()) 86 | matj[l][2*i+1, 2*j+1] = x 87 | matj[l][2*j+1, 2*i+1] = x 88 | x = float(f.readline()) 89 | matj[l][2*i+1,2*i+1] = x 90 | 91 | if l != 2*k1: 92 | k += 1 93 | for i in range(k): 94 | for j in range(k): 95 | x = float(f.readline()) 96 | matj[l][2*k1+2*j+1, 2*i] = x 97 | matj[l][2*i, 2*k1+2*j+1] = x 98 | 99 | if l == 2*k1: 100 | k += 1 101 | else: 102 | k1 += 1 103 | for i in range(k): 104 | for j in range(i): 105 | x = float(f.readline()) 106 | matj[l][2*k1+2*i, 2*k1+2*j] = x 107 | matj[l][2*k1+2*j, 2*k1+2*i] = x 108 | x = float(f.readline()) 109 | matj[l][2*k1+2*i, 2*k1+2*i] = x 110 | 111 | f.close() 112 | return matj 113 | 114 | 115 | def parse_J_single_l(file): 116 | """ 117 | Parse a single numMatJ-XXX.dat file produced by Pinchon & Hoggan's maple script mkNumMatJ.mw 118 | 119 | :param file: 120 | :return: 121 | """ 122 | 123 | f = open(file) 124 | 125 | lval = int(f.readline().split(' ')[1]) 126 | matj = np.zeros((2 * lval + 1, 2 * lval + 1)) 127 | 128 | k = lval / 2 129 | k1 = k 130 | for i in range(k): 131 | for j in range(i): 132 | x = float(f.readline()) 133 | matj[2*i+1, 2*j+1] = x 134 | matj[2*j+1, 2*i+1] = x 135 | x = float(f.readline()) 136 | matj[2*i+1,2*i+1] = x 137 | 138 | if lval != 2 * k1: 139 | k += 1 140 | for i in range(k): 141 | for j in range(k): 142 | x = float(f.readline()) 143 | matj[2*k1+2*j+1, 2*i] = x 144 | matj[2*i, 2*k1+2*j+1] = x 145 | 146 | if lval == 2 * k1: 147 | k += 1 148 | else: 149 | k1 += 1 150 | for i in range(k): 151 | for j in range(i): 152 | x = float(f.readline()) 153 | matj[2*k1+2*i, 2*k1+2*j] = x 154 | matj[2*k1+2*j, 2*k1+2*i] = x 155 | x = float(f.readline()) 156 | matj[2*k1+2*i, 2*k1+2*i] = x 157 | 158 | f.close() 159 | return matj 160 | 161 | 162 | def make_block_J(Jd): 163 | """ 164 | Convert a list of J matrices 0 to N, to block form. 165 | We change the basis on each J matrix (which is assumed to be in the real, quantum-normalized, centered basis, 166 | so that it is in the real, quantum-normalized, block basis. 167 | Then, we extract the blocks. There are 4 blocks for each irrep l, but the middle two are transposes of each other, 168 | so we store only 3 blocks. The outer two blocks are symmetric, but this is not exploited. 169 | 170 | :param Jd: 171 | :return: 172 | """ 173 | Jb = [] 174 | for l in range(len(Jd)): 175 | print('Converting to block matrix. (', l, 'of', len(Jd), ')') 176 | #Bl = c2b(l) 177 | #Jbl = Bl.dot(Jd[l]).dot(Bl.T) 178 | 179 | #c2b = change_of_basis_function(l, 180 | # frm=('real', 'quantum', 'centered', 'cs'), 181 | # to=('real', 'quantum', 'block', 'cs')) 182 | #Jbl = c2b(c2b(Jd[l]).T).T 183 | 184 | Bl = change_of_basis_matrix(l, 185 | frm=('real', 'quantum', 'centered', 'cs'), 186 | to=('real', 'quantum', 'block', 'cs')) 187 | Jbl = Bl.dot(Jd[l]).dot(Bl.T) 188 | 189 | k = l // 2 190 | if l % 2 == 0: 191 | # blocks have dimension k, k, k, k+1 192 | b1 = Jbl[0:k, 0:k] 193 | b2 = Jbl[k:2 * k, 2 * k:3 * k] 194 | #b3 = Jbl[2 * k: 3 * k, k:2 * k] 195 | b4 = Jbl[3 * k:, 3 * k:] 196 | else: 197 | # blocks have dimension k, k+1, k+1, k+1 198 | b1 = Jbl[0:k, 0:k] 199 | b2 = Jbl[k:2 * k + 1, 2 * k + 1:3 * k + 2] 200 | #b3 = Jbl[2 * k + 1:3 * k + 2, k:2 * k + 1] 201 | b4 = Jbl[3 * k + 2:, 3 * k + 2:] 202 | Jb.append([b1, b2, b4]) 203 | return Jb -------------------------------------------------------------------------------- /lie_learn/representations/SO3/spherical_harmonics.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from scipy.special import sph_harm, lpmv 4 | try: 5 | from scipy.misc import factorial 6 | except: 7 | from scipy.special import factorial 8 | 9 | def sh(l, m, theta, phi, field='real', normalization='quantum', condon_shortley=True): 10 | if field == 'real': 11 | return rsh(l, m, theta, phi, normalization, condon_shortley) 12 | elif field == 'complex': 13 | return csh(l, m, theta, phi, normalization, condon_shortley) 14 | else: 15 | raise ValueError('Unknown field: ' + str(field)) 16 | 17 | 18 | def sh_squared_norm(l, normalization='quantum', normalized_haar=True): 19 | """ 20 | Compute the squared norm of the spherical harmonics. 21 | 22 | The squared norm of a function on the sphere is defined as 23 | |f|^2 = int_S^2 |f(x)|^2 dx 24 | where dx is a Haar measure. 25 | 26 | :param l: for some normalization conventions, the norm of a spherical harmonic Y^l_m depends on the degree l 27 | :param normalization: normalization convention for the spherical harmonic 28 | :param normalized_haar: whether to use the Haar measure da db sinb or the normalized Haar measure da db sinb / 4pi 29 | :return: the squared norm of the spherical harmonic with respect to given measure 30 | """ 31 | if normalization == 'quantum' or normalization == 'seismology': 32 | # The quantum and seismology spherical harmonics are normalized with respect to the Haar measure 33 | # dmu(theta, phi) = dtheta sin(theta) dphi 34 | sqnorm = 1. 35 | elif normalization == 'geodesy': 36 | # The geodesy spherical harmonics are normalized with respect to the *normalized* Haar measure 37 | # dmu(theta, phi) = dtheta sin(theta) dphi / 4pi 38 | sqnorm = 4 * np.pi 39 | elif normalization == 'nfft': 40 | sqnorm = 4 * np.pi / (2 * l + 1) 41 | else: 42 | raise ValueError('Unknown normalization') 43 | 44 | if normalized_haar: 45 | return sqnorm / (4 * np.pi) 46 | else: 47 | return sqnorm 48 | 49 | 50 | def block_sh_ph(L_max, theta, phi): 51 | """ 52 | Compute all spherical harmonics up to and including degree L_max, for angles theta and phi. 53 | 54 | This function is currently rather hacky, but the method used here is very fast and stable, compared 55 | to builtin scipy functions. 56 | 57 | :param L_max: 58 | :param theta: 59 | :param phi: 60 | :return: 61 | """ 62 | 63 | from .pinchon_hoggan.pinchon_hoggan import apply_rotation_block, make_c2b 64 | from .irrep_bases import change_of_basis_function 65 | 66 | irreps = np.arange(L_max + 1) 67 | 68 | ls = [[ls] * (2 * ls + 1) for ls in irreps] 69 | ls = np.array([ll for sublist in ls for ll in sublist]) # 0, 1, 1, 1, 2, 2, 2, 2, 2, ... 70 | ms = [list(range(-ls, ls + 1)) for ls in irreps] 71 | ms = np.array([mm for sublist in ms for mm in sublist]) # 0, -1, 0, 1, -2, -1, 0, 1, 2, ... 72 | 73 | # Get a vector Y that selects the 0-frequency component from each irrep in the centered basis 74 | # If D is a Wigner D matrix, then D Y is the center column of D, which is equal to the spherical harmonics. 75 | Y = (ms == 0).astype(float) 76 | 77 | # Change to / from the block basis (since the rotation code works in that basis) 78 | c2b = change_of_basis_function(irreps, 79 | frm=('real', 'quantum', 'centered', 'cs'), 80 | to=('real', 'quantum', 'block', 'cs')) 81 | b2c = change_of_basis_function(irreps, 82 | frm=('real', 'quantum', 'block', 'cs'), 83 | to=('real', 'quantum', 'centered', 'cs')) 84 | 85 | Yb = c2b(Y) 86 | 87 | # Rotate Yb: 88 | c2b = make_c2b(irreps) 89 | import os 90 | J_block = np.load(os.path.join(os.path.dirname(__file__), 'pinchon_hoggan', 'J_block_0-278.npy'), allow_pickle=True) 91 | J_block = list(J_block[irreps]) 92 | 93 | g = np.zeros((theta.size, 3)) 94 | g[:, 0] = phi 95 | g[:, 1] = theta 96 | TYb = apply_rotation_block(g=g, X=Yb[np.newaxis, :], 97 | irreps=irreps, c2b=c2b, 98 | J_block=J_block, l_max=np.max(irreps)) 99 | 100 | print(Yb.shape, TYb.shape) 101 | 102 | # Change back to centered basis 103 | TYc = b2c(TYb.T).T # b2c doesn't work properly for matrices, so do a transpose hack 104 | 105 | print(TYc.shape) 106 | 107 | # Somehow, the SH obtained so far are equal to real, nfft, cs spherical harmonics 108 | # Change to real quantum centered cs 109 | c = change_of_basis_function(irreps, 110 | frm=('real', 'nfft', 'centered', 'cs'), 111 | to=('real', 'quantum', 'centered', 'cs')) 112 | TYc2 = c(TYc) 113 | print(TYc2.shape) 114 | 115 | return TYc2 116 | 117 | 118 | def rsh(l, m, theta, phi, normalization='quantum', condon_shortley=True): 119 | """ 120 | Compute the real spherical harmonic (RSH) S_l^m(theta, phi). 121 | 122 | The RSH are obtained from Complex Spherical Harmonics (CSH) as follows: 123 | if m < 0: 124 | S_l^m = i / sqrt(2) * (Y_l^m - (-1)^m Y_l^{-m}) 125 | if m == 0: 126 | S_l^m = Y_l^0 127 | if m > 0: 128 | S_l^m = 1 / sqrt(2) * (Y_l^{-m} + (-1)^m Y_l^m) 129 | (see [1]) 130 | 131 | Various normalizations for the CSH exist, see the CSH() function. Since the CSH->RSH change of basis is unitary, 132 | the orthogonality and normalization properties of the RSH are the same as those of the CSH from which they were 133 | obtained. Furthermore, the operation of changing normalization and that of changeing field 134 | (complex->real or vice-versa) commute, because the ratio c_m of normalization constants are always the same for 135 | m and -m (to see this that this implies commutativity, substitute Y_l^m * c_m for Y_l^m in the above formula). 136 | 137 | Pinchon & Hoggan [2] define a different change of basis for CSH -> RSH, but they also use an unusual definition 138 | of CSH. To obtain RSH as defined by Pinchon-Hoggan, use this function with normalization='quantum'. 139 | 140 | References: 141 | [1] http://en.wikipedia.org/wiki/Spherical_harmonics#Real_form 142 | [2] Rotation matrices for real spherical harmonics: general rotations of atomic orbitals in space-fixed axes. 143 | 144 | :param l: non-negative integer; the degree of the CSH. 145 | :param m: integer, -l <= m <= l; the order of the CSH. 146 | :param theta: the colatitude / polar angle, 147 | ranging from 0 (North Pole, (X,Y,Z)=(0,0,1)) to pi (South Pole, (X,Y,Z)=(0,0,-1)). 148 | :param phi: the longitude / azimuthal angle, ranging from 0 to 2 pi. 149 | :param normalization: how to normalize the RSH: 150 | 'seismology', 'quantum', 'geodesy'. 151 | these are immediately passed to the CSH functions, and since the change of basis 152 | from CSH to RSH is unitary, the orthogonality and normalization properties are unchanged. 153 | :return: the value of the real spherical harmonic S^l_m(theta, phi) 154 | """ 155 | l, m, theta, phi = np.broadcast_arrays(l, m, theta, phi) 156 | # Get the CSH for m and -m, using Condon-Shortley phase (regardless of whhether CS is requested or not) 157 | # The reason is that the code that changes from CSH to RSH assumes CS phase. 158 | 159 | a = csh(l=l, m=m, theta=theta, phi=phi, normalization=normalization, condon_shortley=True) 160 | b = csh(l=l, m=-m, theta=theta, phi=phi, normalization=normalization, condon_shortley=True) 161 | 162 | #if m > 0: 163 | # y = np.array((b + ((-1.)**m) * a).real / np.sqrt(2.)) 164 | #elif m < 0: 165 | # y = np.array((1j * a - 1j * ((-1.)**(-m)) * b).real / np.sqrt(2.)) 166 | #else: 167 | # # For m == 0, the complex spherical harmonics are already real 168 | # y = np.array(a.real) 169 | 170 | y = ((m > 0) * np.array((b + ((-1.)**m) * a).real / np.sqrt(2.)) 171 | + (m < 0) * np.array((1j * a - 1j * ((-1.)**(-m)) * b).real / np.sqrt(2.)) 172 | + (m == 0) * np.array(a.real)) 173 | 174 | if condon_shortley: 175 | return y 176 | else: 177 | # Cancel the CS phase of y (i.e. multiply by -1 when m is both odd and greater than 0) 178 | return y * ((-1.) ** (m * (m > 0))) 179 | 180 | 181 | def csh(l, m, theta, phi, normalization='quantum', condon_shortley=True): 182 | """ 183 | Compute Complex Spherical Harmonics (CSH) Y_l^m(theta, phi). 184 | Unlike the scipy.special.sph_harm function, we use the common convention that 185 | theta is the polar angle (0 to pi) and phi is the azimuthal angle (0 to 2pi). 186 | 187 | The spherical harmonic 'backbone' is: 188 | Y_l^m(theta, phi) = P_l^m(cos(theta)) exp(i m phi) 189 | where P_l^m is the associated Legendre function as defined in the scipy library (scipy.special.sph_harm). 190 | 191 | Various normalization factors can be multiplied with this function. 192 | -> seismology: sqrt( ((2 l + 1) * (l - m)!) / (4 pi * (l + m)!) ) 193 | -> quantum: (-1)^2 sqrt( ((2 l + 1) * (l - m)!) / (4 pi * (l + m)!) ) 194 | -> unnormalized: 1 195 | -> geodesy: sqrt( ((2 l + 1) * (l - m)!) / (l + m)! ) 196 | -> nfft: sqrt( (l - m)! / (l + m)! ) 197 | 198 | The 'quantum' and 'seismology' CSH are normalized so that 199 | 200 | = 201 | int_S^2 Y_l^m(theta, phi) Y_l'^m'* dOmega 202 | = 203 | delta(l, l') delta(m, m') 204 | where dOmega is the volume element for the sphere S^2: 205 | dOmega = sin(theta) dtheta dphi 206 | The 'geodesy' convention have unit power, meaning the norm is equal to the surface area of the unit sphere (4 pi) 207 | = 4pi delta(l, l') delta(m, m') 208 | So these are orthonormal with respect to the *normalized* Haar measure sin(theta) dtheta dphi / 4pi 209 | 210 | On each of these normalizations, one can optionally include a Condon-Shortley phase factor: 211 | (-1)^m (if m > 0) 212 | 1 (otherwise) 213 | Note that this is the definition of Condon-Shortley according to wikipedia [1], but other sources call a 214 | phase factor of (-1)^m a Condon-Shortley phase (without mentioning the condition m > 0). 215 | 216 | References: 217 | [1] http://en.wikipedia.org/wiki/Spherical_harmonics#Conventions 218 | 219 | :param l: non-negative integer; the degree of the CSH. 220 | :param m: integer, -l <= m <= l; the order of the CSH. 221 | :param theta: the colatitude / polar angle, 222 | ranging from 0 (North Pole, (X,Y,Z)=(0,0,1)) to pi (South Pole, (X,Y,Z)=(0,0,-1)). 223 | :param phi: the longitude / azimuthal angle, ranging from 0 to 2 pi. 224 | :param normalization: how to normalize the CSH: 225 | 'seismology', 'quantum', 'geodesy', 'unnormalized', 'nfft'. 226 | :return: the value of the complex spherical harmonic Y^l_m(theta, phi) 227 | """ 228 | # NOTE: it seems like in the current version of scipy.special, sph_harm no longer accepts keyword arguments, 229 | # so I'm removing them. I hope the order of args hasn't changed 230 | if normalization == 'quantum': 231 | # y = ((-1.) ** m) * sph_harm(m, l, theta=phi, phi=theta) 232 | y = ((-1.) ** m) * sph_harm(m, l, phi, theta) 233 | elif normalization == 'seismology': 234 | # y = sph_harm(m, l, theta=phi, phi=theta) 235 | y = sph_harm(m, l, phi, theta) 236 | elif normalization == 'geodesy': 237 | # y = np.sqrt(4 * np.pi) * sph_harm(m, l, theta=phi, phi=theta) 238 | y = np.sqrt(4 * np.pi) * sph_harm(m, l, phi, theta) 239 | elif normalization == 'unnormalized': 240 | # y = sph_harm(m, l, theta=phi, phi=theta) / np.sqrt((2 * l + 1) * factorial(l - m) / 241 | # (4 * np.pi * factorial(l + m))) 242 | y = sph_harm(m, l, phi, theta) / np.sqrt((2 * l + 1) * factorial(l - m) / 243 | (4 * np.pi * factorial(l + m))) 244 | elif normalization == 'nfft': 245 | # y = sph_harm(m, l, theta=phi, phi=theta) / np.sqrt((2 * l + 1) / (4 * np.pi)) 246 | y = sph_harm(m, l, phi, theta) / np.sqrt((2 * l + 1) / (4 * np.pi)) 247 | else: 248 | raise ValueError('Unknown normalization convention:' + str(normalization)) 249 | 250 | if condon_shortley: 251 | # The sph_harm function already includes CS phase 252 | return y 253 | else: 254 | # Cancel the CS phase in sph_harm (i.e. multiply by -1 when m is both odd and greater than 0) 255 | return y * ((-1.) ** (m * (m > 0))) 256 | 257 | 258 | # For testing only: 259 | def _naive_csh_unnormalized(l, m, theta, phi): 260 | """ 261 | Compute unnormalized SH 262 | """ 263 | return lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi) 264 | 265 | 266 | def _naive_csh_quantum(l, m, theta, phi): 267 | """ 268 | Compute orthonormalized spherical harmonics in a naive way. 269 | """ 270 | return (((-1.) ** m) * lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi) * 271 | np.sqrt(((2 * l + 1) * factorial(l - m)) 272 | / 273 | (4 * np.pi * factorial(l + m)))) 274 | 275 | 276 | def _naive_csh_seismology(l, m, theta, phi): 277 | """ 278 | Compute the spherical harmonics according to the seismology convention, in a naive way. 279 | This appears to be equal to the sph_harm function in scipy.special. 280 | """ 281 | return (lpmv(m, l, np.cos(theta)) * np.exp(1j * m * phi) * 282 | np.sqrt(((2 * l + 1) * factorial(l - m)) 283 | / 284 | (4 * np.pi * factorial(l + m)))) 285 | 286 | 287 | def _naive_csh_ph(l, m, theta, phi): 288 | """ 289 | CSH as defined by Pinchon-Hoggan. Same as wikipedia's quantum-normalized SH = naive_Y_quantum() 290 | """ 291 | if l == 0 and m == 0: 292 | return 1. / np.sqrt(4 * np.pi) 293 | else: 294 | phase = ((1j) ** (m + np.abs(m))) 295 | normalizer = np.sqrt(((2 * l + 1.) * factorial(l - np.abs(m))) 296 | / 297 | (4 * np.pi * factorial(l + np.abs(m)))) 298 | P = lpmv(np.abs(m), l, np.cos(theta)) 299 | e = np.exp(1j * m * phi) 300 | return phase * normalizer * P * e 301 | 302 | 303 | def _naive_rsh_ph(l, m, theta, phi): 304 | 305 | if m == 0: 306 | return np.sqrt((2 * l + 1.) / (4 * np.pi)) * lpmv(m, l, np.cos(theta)) 307 | elif m < 0: 308 | return np.sqrt(((2 * l + 1.) * factorial(l + m)) / 309 | (2 * np.pi * factorial(l - m))) * lpmv(-m, l, np.cos(theta)) * np.sin(-m * phi) 310 | elif m > 0: 311 | return np.sqrt(((2 * l + 1.) * factorial(l - m)) / 312 | (2 * np.pi * factorial(l + m))) * lpmv(m, l, np.cos(theta)) * np.cos(m * phi) 313 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/test_SO3_irrep_bases.py: -------------------------------------------------------------------------------- 1 | from lie_learn.representations.SO3.irrep_bases import * 2 | from .spherical_harmonics import * 3 | 4 | TEST_L_MAX = 5 5 | 6 | def test_change_of_basis_matrix(): 7 | """ 8 | Testing if change of basis matrix is consistent with spherical harmonics functions 9 | """ 10 | 11 | for l in range(TEST_L_MAX): 12 | theta = np.random.rand() * np.pi 13 | phi = np.random.rand() * np.pi * 2 14 | for from_field in ['complex', 'real']: 15 | for from_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 16 | for from_cs in ['cs', 'nocs']: 17 | for to_field in ['complex', 'real']: 18 | for to_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 19 | for to_cs in ['cs', 'nocs']: 20 | Y_from = sh(l, np.arange(-l, l + 1), theta, phi, 21 | from_field, from_normalization, from_cs == 'cs') 22 | 23 | Y_to = sh(l, np.arange(-l, l + 1), theta, phi, 24 | to_field, to_normalization, to_cs == 'cs') 25 | 26 | B = change_of_basis_matrix(l=l, 27 | frm=(from_field, from_normalization, 'centered', from_cs), 28 | to=(to_field, to_normalization, 'centered', to_cs)) 29 | 30 | print(from_field, from_normalization, from_cs, '->', to_field, to_normalization, to_cs, np.sum(np.abs(B.dot(Y_from) - Y_to))) 31 | assert np.isclose(np.sum(np.abs(B.dot(Y_from) - Y_to)), 0.0) 32 | assert np.isclose(np.sum(np.abs(np.linalg.inv(B).dot(Y_to) - Y_from)), 0.0) 33 | 34 | 35 | def test_change_of_basis_function(): 36 | """ 37 | Testing if change of basis function is consistent with spherical harmonics functions 38 | """ 39 | 40 | for l in range(TEST_L_MAX): 41 | theta = np.random.rand() * np.pi 42 | phi = np.random.rand() * np.pi * 2 43 | for from_field in ['complex', 'real']: 44 | for from_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 45 | for from_cs in ['cs', 'nocs']: 46 | for to_field in ['complex', 'real']: 47 | for to_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 48 | for to_cs in ['cs', 'nocs']: 49 | 50 | Y_from = sh(l, np.arange(-l, l + 1), theta, phi, 51 | from_field, from_normalization, from_cs == 'cs') 52 | 53 | Y_to = sh(l, np.arange(-l, l + 1), theta, phi, 54 | to_field, to_normalization, to_cs == 'cs') 55 | 56 | f = change_of_basis_function(l=l, 57 | frm=(from_field, from_normalization, 'centered', from_cs), 58 | to=(to_field, to_normalization, 'centered', to_cs)) 59 | 60 | print(from_field, from_normalization, from_cs, '->', to_field, to_normalization, to_cs, np.sum(np.abs(f(Y_from) - Y_to))) 61 | assert np.isclose(np.sum(np.abs(f(Y_from) - Y_to)), 0.0) 62 | 63 | 64 | def test_change_of_basis_function_lists(): 65 | """ 66 | Testing change of basis function for spherical harmonics for multiple orders at once. 67 | The change-of-basis function for spherical harmonics should be consistent with the CSH & RSH functions. 68 | """ 69 | l = np.arange(4) 70 | ls = np.array([0, 1,1,1, 2,2,2,2,2, 3,3,3,3,3,3,3]) 71 | ms = np.array([0, -1,0,1, -2,-1,0,1,2, -3,-2,-1,0,1,2,3]) 72 | 73 | theta = np.random.rand() * np.pi 74 | phi = np.random.rand() * np.pi * 2 75 | for from_field in ['complex', 'real']: 76 | for from_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 77 | for from_cs in ['cs', 'nocs']: 78 | for to_field in ['complex', 'real']: 79 | for to_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 80 | for to_cs in ['cs', 'nocs']: 81 | 82 | Y_from = sh(ls, ms, theta, phi, 83 | from_field, from_normalization, from_cs == 'cs') 84 | 85 | Y_to = sh(ls, ms, theta, phi, 86 | to_field, to_normalization, to_cs == 'cs') 87 | 88 | f = change_of_basis_function(l=l, 89 | frm=(from_field, from_normalization, 'centered', from_cs), 90 | to=(to_field, to_normalization, 'centered', to_cs)) 91 | 92 | print(from_field, from_normalization, from_cs, '->', to_field, to_normalization, to_cs, np.sum(np.abs(f(Y_from) - Y_to))) 93 | assert np.isclose(np.sum(np.abs(f(Y_from) - Y_to)), 0.0) 94 | 95 | 96 | def test_invertibility(): 97 | """ 98 | Testing if change_of_basis_function for SO(3) is invertible 99 | """ 100 | 101 | for l in range(TEST_L_MAX): 102 | theta = np.random.rand() * np.pi 103 | phi = np.random.rand() * np.pi * 2 104 | for from_field in ['complex', 'real']: 105 | for from_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 106 | for from_cs in ['cs', 'nocs']: 107 | for from_order in ['centered', 'block']: 108 | for to_field in ['complex', 'real']: 109 | for to_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 110 | for to_cs in ['cs', 'nocs']: 111 | for to_order in ['centered', 'block']: 112 | # A truly complex function cannot be made real; 113 | if from_field == 'complex' and to_field == 'real': 114 | continue 115 | 116 | if from_field == 'complex': 117 | Y = np.random.randn(2 * l + 1) + np.random.randn(2 * l + 1) * 1j 118 | else: 119 | Y = np.random.randn(2 * l + 1) 120 | 121 | f = change_of_basis_function(l=l, 122 | frm=(from_field, from_normalization, from_order, from_cs), 123 | to=(to_field, to_normalization, to_order, to_cs)) 124 | 125 | f_inv = change_of_basis_function(l=l, 126 | frm=(to_field, to_normalization, to_order, to_cs), 127 | to=(from_field, from_normalization, from_order, from_cs)) 128 | 129 | 130 | print(from_field, from_normalization, from_cs, from_order, '->', to_field, to_normalization, to_cs, to_order, np.sum(np.abs(f_inv(f(Y)) - Y))) 131 | assert np.isclose(np.sum(np.abs(f_inv(f(Y)) - Y)), 0.) 132 | #assert np.isclose(np.sum(np.abs(f(f_inv(Y)) - Y)), 0.) 133 | 134 | 135 | def test_linearity_change_of_basis(): 136 | """ 137 | Testing that SO3 change of basis is indeed linear 138 | """ 139 | for l in range(TEST_L_MAX): 140 | theta = np.random.rand() * np.pi 141 | phi = np.random.rand() * np.pi * 2 142 | for from_field in ['complex', 'real']: 143 | for from_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 144 | for from_cs in ['cs', 'nocs']: 145 | for from_order in ['centered', 'block']: 146 | for to_field in ['complex', 'real']: 147 | for to_normalization in ['seismology', 'quantum', 'geodesy', 'nfft']: 148 | for to_cs in ['cs', 'nocs']: 149 | for to_order in ['centered', 'block']: 150 | 151 | # A truly complex function cannot be made real; 152 | if from_field == 'complex' and to_field == 'real': 153 | continue 154 | 155 | Y1 = np.random.randn(2 * l + 1) 156 | Y2 = np.random.randn(2 * l + 1) 157 | a = np.random.randn(1) 158 | b = np.random.randn(1) 159 | 160 | f = change_of_basis_function(l=l, 161 | frm=(from_field, from_normalization, from_order, from_cs), 162 | 163 | to=(to_field, to_normalization, from_order, to_cs)) 164 | 165 | print(from_field, from_normalization, from_cs, from_order, '->', to_field, to_normalization, to_cs, to_order, np.sum(np.abs(a * f(Y1) + b * f(Y2) - f(a*Y1 + b*Y2)))) 166 | assert np.isclose(np.sum(np.abs(a * f(Y1) + b * f(Y2) - f(a*Y1 + b*Y2))), 0.) 167 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/test_spherical_harmonics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import lie_learn.spaces.S2 as S2 3 | from lie_learn.representations.SO3.spherical_harmonics import sh, sh_squared_norm 4 | 5 | 6 | def check_orthogonality(L_max=3, grid_type='Gauss-Legendre', 7 | field='real', normalization='quantum', condon_shortley=True): 8 | 9 | theta, phi = S2.meshgrid(b=L_max + 1, grid_type=grid_type) 10 | w = S2.quadrature_weights(b=L_max + 1, grid_type=grid_type) 11 | 12 | for l in range(L_max): 13 | for m in range(-l, l + 1): 14 | for l2 in range(L_max): 15 | for m2 in range(-l2, l2 + 1): 16 | Ylm = sh(l, m, theta, phi, field, normalization, condon_shortley) 17 | Ylm2 = sh(l2, m2, theta, phi, field, normalization, condon_shortley) 18 | 19 | dot_numerical = S2.integrate_quad(Ylm * Ylm2.conj(), grid_type=grid_type, normalize=False, w=w) 20 | 21 | dot_numerical2 = S2.integrate( 22 | lambda t, p: sh(l, m, t, p, field, normalization, condon_shortley) * \ 23 | sh(l2, m2, t, p, field, normalization, condon_shortley).conj(), normalize=False) 24 | 25 | sqnorm_analytical = sh_squared_norm(l, normalization, normalized_haar=False) 26 | dot_analytical = sqnorm_analytical * (l == l2 and m == m2) 27 | 28 | print(l, m, l2, m2, field, normalization, condon_shortley, dot_analytical, dot_numerical, dot_numerical2) 29 | assert np.isclose(dot_numerical, dot_analytical) 30 | assert np.isclose(dot_numerical2, dot_analytical) 31 | 32 | 33 | def test_orthogonality(): 34 | L_max = 2 35 | grid_type = 'Gauss-Legendre' 36 | 37 | for field in ('real', 'complex'): 38 | for normalization in ('quantum', 'seismology', 'geodesy', 'nfft'): 39 | for condon_shortley in (True, False): 40 | check_orthogonality(L_max, grid_type, field, normalization, condon_shortley) 41 | -------------------------------------------------------------------------------- /lie_learn/representations/SO3/wigner_d.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | from lie_learn.representations.SO3.pinchon_hoggan.pinchon_hoggan_dense import Jd, rot_mat 5 | from lie_learn.representations.SO3.irrep_bases import change_of_basis_matrix 6 | 7 | 8 | def wigner_d_matrix(l, beta, 9 | field='real', normalization='quantum', order='centered', condon_shortley='cs'): 10 | """ 11 | Compute the Wigner-d matrix of degree l at beta, in the basis defined by 12 | (field, normalization, order, condon_shortley) 13 | 14 | The Wigner-d matrix of degree l has shape (2l + 1) x (2l + 1). 15 | 16 | :param l: the degree of the Wigner-d function. l >= 0 17 | :param beta: the argument. 0 <= beta <= pi 18 | :param field: 'real' or 'complex' 19 | :param normalization: 'quantum', 'seismology', 'geodesy' or 'nfft' 20 | :param order: 'centered' or 'block' 21 | :param condon_shortley: 'cs' or 'nocs' 22 | :return: d^l_mn(beta) in the chosen basis 23 | """ 24 | # This returns the d matrix in the (real, quantum-normalized, centered, cs) convention 25 | d = rot_mat(alpha=0., beta=beta, gamma=0., l=l, J=Jd[l]) 26 | 27 | if (field, normalization, order, condon_shortley) != ('real', 'quantum', 'centered', 'cs'): 28 | # TODO use change of basis function instead of matrix? 29 | B = change_of_basis_matrix( 30 | l, 31 | frm=('real', 'quantum', 'centered', 'cs'), 32 | to=(field, normalization, order, condon_shortley)) 33 | BB = change_of_basis_matrix( 34 | l, 35 | frm=(field, normalization, order, condon_shortley), 36 | to=('real', 'quantum', 'centered', 'cs')) 37 | d = B.dot(d).dot(BB) 38 | 39 | # The Wigner-d matrices are always real, even in the complex basis 40 | # (I tested this numerically, and have seen it in several texts) 41 | # assert np.isclose(np.sum(np.abs(d.imag)), 0.0) 42 | d = d.real 43 | 44 | return d 45 | 46 | 47 | def wigner_D_matrix(l, alpha, beta, gamma, 48 | field='real', normalization='quantum', order='centered', condon_shortley='cs'): 49 | """ 50 | Evaluate the Wigner-d matrix D^l_mn(alpha, beta, gamma) 51 | 52 | :param l: the degree of the Wigner-d function. l >= 0 53 | :param alpha: the argument. 0 <= alpha <= 2 pi 54 | :param beta: the argument. 0 <= beta <= pi 55 | :param gamma: the argument. 0 <= gamma <= 2 pi 56 | :param field: 'real' or 'complex' 57 | :param normalization: 'quantum', 'seismology', 'geodesy' or 'nfft' 58 | :param order: 'centered' or 'block' 59 | :param condon_shortley: 'cs' or 'nocs' 60 | :return: D^l_mn(alpha, beta, gamma) in the chosen basis 61 | """ 62 | 63 | D = rot_mat(alpha=alpha, beta=beta, gamma=gamma, l=l, J=Jd[l]) 64 | 65 | if (field, normalization, order, condon_shortley) != ('real', 'quantum', 'centered', 'cs'): 66 | B = change_of_basis_matrix( 67 | l, 68 | frm=('real', 'quantum', 'centered', 'cs'), 69 | to=(field, normalization, order, condon_shortley)) 70 | BB = change_of_basis_matrix( 71 | l, 72 | frm=(field, normalization, order, condon_shortley), 73 | to=('real', 'quantum', 'centered', 'cs')) 74 | D = B.dot(D).dot(BB) 75 | 76 | if field == 'real': 77 | # print('WIGNER D IMAG PART:', np.sum(np.abs(D.imag))) 78 | assert np.isclose(np.sum(np.abs(D.imag)), 0.0) 79 | D = D.real 80 | 81 | return D 82 | 83 | 84 | def wigner_d_function(l, m, n, beta, 85 | field='real', normalization='quantum', order='centered', condon_shortley='cs'): 86 | """ 87 | Evaluate a single Wigner-d function d^l_mn(beta) 88 | 89 | NOTE: for now, we implement this by computing the entire degree-l Wigner-d matrix and then selecting 90 | the (m,n) element, so this function is not fast. 91 | 92 | :param l: the degree of the Wigner-d function. l >= 0 93 | :param m: the order of the Wigner-d function. -l <= m <= l 94 | :param n: the order of the Wigner-d function. -l <= n <= l 95 | :param beta: the argument. 0 <= beta <= pi 96 | :param field: 'real' or 'complex' 97 | :param normalization: 'quantum', 'seismology', 'geodesy' or 'nfft' 98 | :param order: 'centered' or 'block' 99 | :param condon_shortley: 'cs' or 'nocs' 100 | :return: d^l_mn(beta) in the chosen basis 101 | """ 102 | return wigner_d_matrix(l, beta, field, normalization, order, condon_shortley)[l + m, l + n] 103 | 104 | 105 | def wigner_D_function(l, m, n, alpha, beta, gamma, 106 | field='real', normalization='quantum', order='centered', condon_shortley='cs'): 107 | """ 108 | Evaluate a single Wigner-d function d^l_mn(beta) 109 | 110 | NOTE: for now, we implement this by computing the entire degree-l Wigner-D matrix and then selecting 111 | the (m,n) element, so this function is not fast. 112 | 113 | :param l: the degree of the Wigner-d function. l >= 0 114 | :param m: the order of the Wigner-d function. -l <= m <= l 115 | :param n: the order of the Wigner-d function. -l <= n <= l 116 | :param alpha: the argument. 0 <= alpha <= 2 pi 117 | :param beta: the argument. 0 <= beta <= pi 118 | :param gamma: the argument. 0 <= gamma <= 2 pi 119 | :param field: 'real' or 'complex' 120 | :param normalization: 'quantum', 'seismology', 'geodesy' or 'nfft' 121 | :param order: 'centered' or 'block' 122 | :param condon_shortley: 'cs' or 'nocs' 123 | :return: d^l_mn(beta) in the chosen basis 124 | """ 125 | return wigner_D_matrix(l, alpha, beta, gamma, field, normalization, order, condon_shortley)[l + m, l + n] 126 | 127 | 128 | def wigner_D_norm(l, normalized_haar=True): 129 | """ 130 | Compute the squared norm of the Wigner-D functions. 131 | 132 | The squared norm of a function on the SO(3) is defined as 133 | |f|^2 = int_SO(3) |f(g)|^2 dg 134 | where dg is a Haar measure. 135 | 136 | :param l: for some normalization conventions, the norm of a Wigner-D function D^l_mn depends on the degree l 137 | :param normalized_haar: whether to use the Haar measure da db sinb dc or the normalized Haar measure 138 | da db sinb dc / 8pi^2 139 | :return: the squared norm of the spherical harmonic with respect to given measure 140 | 141 | :param l: 142 | :param normalization: 143 | :return: 144 | """ 145 | if normalized_haar: 146 | return 1. / (2 * l + 1) 147 | else: 148 | return (8 * np.pi ** 2) / (2 * l + 1) 149 | 150 | 151 | def wigner_d_naive(l, m, n, beta): 152 | """ 153 | Numerically naive implementation of the Wigner-d function. 154 | This is useful for checking the correctness of other implementations. 155 | 156 | :param l: the degree of the Wigner-d function. l >= 0 157 | :param m: the order of the Wigner-d function. -l <= m <= l 158 | :param n: the order of the Wigner-d function. -l <= n <= l 159 | :param beta: the argument. 0 <= beta <= pi 160 | :return: d^l_mn(beta) in the TODO: what basis? complex, quantum(?), centered, cs(?) 161 | """ 162 | from scipy.special import eval_jacobi 163 | try: 164 | from scipy.misc import factorial 165 | except: 166 | from scipy.special import factorial 167 | 168 | from sympy.functions.special.polynomials import jacobi, jacobi_normalized 169 | from sympy.abc import j, a, b, x 170 | from sympy import N 171 | #jfun = jacobi_normalized(j, a, b, x) 172 | jfun = jacobi(j, a, b, x) 173 | # eval_jacobi = lambda q, r, p, o: float(jfun.eval(int(q), int(r), int(p), float(o))) 174 | # eval_jacobi = lambda q, r, p, o: float(N(jfun, int(q), int(r), int(p), float(o))) 175 | eval_jacobi = lambda q, r, p, o: float(jfun.subs({j:int(q), a:int(r), b:int(p), x:float(o)})) 176 | 177 | mu = np.abs(m - n) 178 | nu = np.abs(m + n) 179 | s = l - (mu + nu) / 2 180 | xi = 1 if n >= m else (-1) ** (n - m) 181 | 182 | # print(s, mu, nu, np.cos(beta), type(s), type(mu), type(nu), type(np.cos(beta))) 183 | jac = eval_jacobi(s, mu, nu, np.cos(beta)) 184 | z = np.sqrt((factorial(s) * factorial(s + mu + nu)) / (factorial(s + mu) * factorial(s + nu))) 185 | 186 | # print(l, m, n, beta, np.isfinite(mu), np.isfinite(nu), np.isfinite(s), np.isfinite(xi), np.isfinite(jac), np.isfinite(z)) 187 | assert np.isfinite(mu) and np.isfinite(nu) and np.isfinite(s) and np.isfinite(xi) and np.isfinite(jac) and np.isfinite(z) 188 | assert np.isfinite(xi * z * np.sin(beta / 2) ** mu * np.cos(beta / 2) ** nu * jac) 189 | return xi * z * np.sin(beta / 2) ** mu * np.cos(beta / 2) ** nu * jac 190 | 191 | 192 | def wigner_d_naive_v2(l, m, n, beta): 193 | """ 194 | Wigner d functions as defined in the SOFT 2.0 documentation. 195 | When approx_lim is set to a high value, this function appears to give 196 | identical results to Johann Goetz' wignerd() function. 197 | 198 | However, integration fails: does not satisfy orthogonality relations everywhere... 199 | """ 200 | from scipy.special import jacobi 201 | 202 | if n >= m: 203 | xi = 1 204 | else: 205 | xi = (-1)**(n - m) 206 | 207 | mu = np.abs(m - n) 208 | nu = np.abs(n + m) 209 | s = l - (mu + nu) * 0.5 210 | 211 | sq = np.sqrt((np.math.factorial(s) * np.math.factorial(s + mu + nu)) 212 | / (np.math.factorial(s + mu) * np.math.factorial(s + nu))) 213 | sinb = np.sin(beta * 0.5) ** mu 214 | cosb = np.cos(beta * 0.5) ** nu 215 | P = jacobi(s, mu, nu)(np.cos(beta)) 216 | return xi * sq * sinb * cosb * P 217 | 218 | 219 | def wigner_d_naive_v3(l, m, n, approx_lim=1000000): 220 | """ 221 | Wigner "small d" matrix. (Euler z-y-z convention) 222 | example: 223 | l = 2 224 | m = 1 225 | n = 0 226 | beta = linspace(0,pi,100) 227 | wd210 = wignerd(l,m,n)(beta) 228 | 229 | some conditions have to be met: 230 | l >= 0 231 | -l <= m <= l 232 | -l <= n <= l 233 | 234 | The approx_lim determines at what point 235 | bessel functions are used. Default is when: 236 | l > m+10 237 | and 238 | l > n+10 239 | 240 | for integer l and n=0, we can use the spherical harmonics. If in 241 | addition m=0, we can use the ordinary legendre polynomials. 242 | """ 243 | from scipy.special import jv, legendre, sph_harm, jacobi 244 | try: 245 | from scipy.misc import factorial, comb 246 | except: 247 | from scipy.special import factorial, comb 248 | from numpy import floor, sqrt, sin, cos, exp, power 249 | from math import pi 250 | from scipy.special import jacobi 251 | 252 | if (l < 0) or (abs(m) > l) or (abs(n) > l): 253 | raise ValueError("wignerd(l = {0}, m = {1}, n = {2}) value error.".format(l, m, n) \ 254 | + " Valid range for parameters: l>=0, -l<=m,n<=l.") 255 | 256 | if (l > (m + approx_lim)) and (l > (n + approx_lim)): 257 | #print 'bessel (approximation)' 258 | return lambda beta: jv(m - n, l * beta) 259 | 260 | if (floor(l) == l) and (n == 0): 261 | if m == 0: 262 | #print 'legendre (exact)' 263 | return lambda beta: legendre(l)(cos(beta)) 264 | elif False: 265 | #print 'spherical harmonics (exact)' 266 | a = sqrt(4. * pi / (2. * l + 1.)) 267 | return lambda beta: a * sph_harm(m, l, beta, 0.).conj() 268 | 269 | jmn_terms = { 270 | l + n : (m - n, m - n), 271 | l - n : (n - m, 0.), 272 | l + m : (n - m, 0.), 273 | l - m : (m - n, m - n), 274 | } 275 | 276 | k = min(jmn_terms) 277 | a, lmb = jmn_terms[k] 278 | 279 | b = 2. * l - 2. * k - a 280 | 281 | if (a < 0) or (b < 0): 282 | raise ValueError("wignerd(l = {0}, m = {1}, n = {2}) value error.".format(l, m, n) \ 283 | + " Encountered negative values in (a,b) = ({0},{1})".format(a,b)) 284 | 285 | coeff = power(-1.,lmb) * sqrt(comb(2. * l - k, k + a)) * (1. / sqrt(comb(k + b, b))) 286 | 287 | #print 'jacobi (exact)' 288 | return lambda beta: coeff \ 289 | * power(sin(0.5*beta),a) \ 290 | * power(cos(0.5*beta),b) \ 291 | * jacobi(k,a,b)(cos(beta)) 292 | -------------------------------------------------------------------------------- /lie_learn/representations/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/representations/__init__.py -------------------------------------------------------------------------------- /lie_learn/spaces/S2.py: -------------------------------------------------------------------------------- 1 | """ 2 | The 2-sphere, S^2 3 | """ 4 | import numpy as np 5 | from numpy.polynomial.legendre import leggauss 6 | 7 | 8 | def change_coordinates(coords, p_from='C', p_to='S'): 9 | """ 10 | Change Spherical to Cartesian coordinates and vice versa, for points x in S^2. 11 | 12 | In the spherical system, we have coordinates beta and alpha, 13 | where beta in [0, pi] and alpha in [0, 2pi] 14 | 15 | We use the names beta and alpha for compatibility with the SO(3) code (S^2 being a quotient SO(3)/SO(2)). 16 | Many sources, like wikipedia use theta=beta and phi=alpha. 17 | 18 | :param coords: coordinate array 19 | :param p_from: 'C' for Cartesian or 'S' for spherical coordinates 20 | :param p_to: 'C' for Cartesian or 'S' for spherical coordinates 21 | :return: new coordinates 22 | """ 23 | if p_from == p_to: 24 | return coords 25 | elif p_from == 'S' and p_to == 'C': 26 | 27 | beta = coords[..., 0] 28 | alpha = coords[..., 1] 29 | r = 1. 30 | 31 | out = np.empty(beta.shape + (3,)) 32 | 33 | ct = np.cos(beta) 34 | cp = np.cos(alpha) 35 | st = np.sin(beta) 36 | sp = np.sin(alpha) 37 | out[..., 0] = r * st * cp # x 38 | out[..., 1] = r * st * sp # y 39 | out[..., 2] = r * ct # z 40 | return out 41 | 42 | elif p_from == 'C' and p_to == 'S': 43 | 44 | x = coords[..., 0] 45 | y = coords[..., 1] 46 | z = coords[..., 2] 47 | 48 | out = np.empty(x.shape + (2,)) 49 | out[..., 0] = np.arccos(z) # beta 50 | out[..., 1] = np.arctan2(y, x) # alpha 51 | return out 52 | 53 | else: 54 | raise ValueError('Unknown conversion:' + str(p_from) + ' to ' + str(p_to)) 55 | 56 | 57 | def meshgrid(b, grid_type='Driscoll-Healy'): 58 | """ 59 | Create a coordinate grid for the 2-sphere. 60 | There are various ways to setup a grid on the sphere. 61 | 62 | if grid_type == 'Driscoll-Healy', we follow the grid_type from [4], which is also used in [5]: 63 | beta_j = pi j / (2 b) for j = 0, ..., 2b - 1 64 | alpha_k = pi k / b for k = 0, ..., 2b - 1 65 | 66 | if grid_type == 'SOFT', we follow the grid_type from [1] and [6] 67 | beta_j = pi (2 j + 1) / (4 b) for j = 0, ..., 2b - 1 68 | alpha_k = pi k / b for k = 0, ..., 2b - 1 69 | 70 | if grid_type == 'Clenshaw-Curtis', we use the Clenshaw-Curtis grid, as defined in [2] (section 6): 71 | beta_j = j pi / (2b) for j = 0, ..., 2b 72 | alpha_k = k pi / (b + 1) for k = 0, ..., 2b + 1 73 | 74 | if grid_type == 'Gauss-Legendre', we use the Gauss-Legendre grid, as defined in [2] (section 6) and [7] (eq. 2): 75 | beta_j = the Gauss-Legendre nodes for j = 0, ..., b 76 | alpha_k = k pi / (b + 1), for k = 0, ..., 2 b + 1 77 | 78 | if grid_type == 'HEALPix', we use the HEALPix grid, see [2] (section 6): 79 | TODO 80 | 81 | if grid_type == 'equidistribution', we use the equidistribution grid, as defined in [2] (section 6): 82 | TODO 83 | 84 | [1] SOFT: SO(3) Fourier Transforms 85 | Kostelec, Peter J & Rockmore, Daniel N. 86 | 87 | [2] Fast evaluation of quadrature formulae on the sphere 88 | Jens Keiner, Daniel Potts 89 | 90 | [3] A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature 91 | Zydrunas Gimbutas Shravan Veerapaneni 92 | 93 | [4] Computing Fourier transforms and convolutions on the 2-sphere 94 | Driscoll, JR & Healy, DM 95 | 96 | [5] Engineering Applications of Noncommutative Harmonic Analysis 97 | Chrikjian, G.S. & Kyatkin, A.B. 98 | 99 | [6] FFTs for the 2-Sphere – Improvements and Variations 100 | Healy, D., Rockmore, D., Kostelec, P., Moore, S 101 | 102 | [7] A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature 103 | Zydrunas Gimbutas, Shravan Veerapaneni 104 | 105 | :param b: the bandwidth / resolution 106 | :return: a meshgrid on S^2 107 | """ 108 | return np.meshgrid(*linspace(b, grid_type), indexing='ij') 109 | 110 | 111 | def linspace(b, grid_type='Driscoll-Healy'): 112 | if grid_type == 'Driscoll-Healy': 113 | beta = np.arange(2 * b) * np.pi / (2. * b) 114 | alpha = np.arange(2 * b) * np.pi / b 115 | elif grid_type == 'SOFT': 116 | beta = np.pi * (2 * np.arange(2 * b) + 1) / (4. * b) 117 | alpha = np.arange(2 * b) * np.pi / b 118 | elif grid_type == 'Clenshaw-Curtis': 119 | # beta = np.arange(2 * b + 1) * np.pi / (2 * b) 120 | # alpha = np.arange(2 * b + 2) * np.pi / (b + 1) 121 | # Must use np.linspace to prevent numerical errors that cause beta > pi 122 | beta = np.linspace(0, np.pi, 2 * b + 1) 123 | alpha = np.linspace(0, 2 * np.pi, 2 * b + 2, endpoint=False) 124 | elif grid_type == 'Gauss-Legendre': 125 | x, _ = leggauss(b + 1) # TODO: leggauss docs state that this may not be only stable for orders > 100 126 | beta = np.arccos(x) 127 | alpha = np.arange(2 * b + 2) * np.pi / (b + 1) 128 | elif grid_type == 'HEALPix': 129 | #TODO: implement this here so that we don't need the dependency on healpy / healpix_compat 130 | from healpix_compat import healpy_sphere_meshgrid 131 | return healpy_sphere_meshgrid(b) 132 | elif grid_type == 'equidistribution': 133 | raise NotImplementedError('Not implemented yet; see Fast evaluation of quadrature formulae on the sphere.') 134 | else: 135 | raise ValueError('Unknown grid_type:' + grid_type) 136 | return beta, alpha 137 | 138 | 139 | def quadrature_weights(b, grid_type='Gauss-Legendre'): 140 | """ 141 | Compute quadrature weights for a given grid-type. 142 | The function S2.meshgrid generates the points that correspond to the weights generated by this function. 143 | 144 | if convention == 'Gauss-Legendre': 145 | The quadrature formula is exact for polynomials up to degree M less than or equal to 2b + 1, 146 | so that we can compute exact Fourier coefficients for f a polynomial of degree at most b. 147 | 148 | if convention == 'Clenshaw-Curtis': 149 | The quadrature formula is exact for polynomials up to degree M less than or equal to 2b, 150 | so that we can compute exact Fourier coefficients for f a polynomial of degree at most b. 151 | 152 | :param b: the grid resolution. See S2.meshgrid 153 | :param grid_type: 154 | :return: 155 | """ 156 | if grid_type == 'Clenshaw-Curtis': 157 | # There is a faster fft based method to compute these weights 158 | # see "Fast evaluation of quadrature formulae on the sphere" 159 | # W = np.empty((2 * b + 2, 2 * b + 1)) 160 | # for j in range(2 * b + 1): 161 | # eps_j_2b = 0.5 if j == 0 or j == 2 * b else 1. 162 | # for k in range(2 * b + 2): # Doesn't seem to depend on k.. 163 | # W[k, j] = (4 * np.pi * eps_j_2b) / (b * (2 * b + 2)) 164 | # sum = 0. 165 | # for l in range(b + 1): 166 | # eps_l_b = 0.5 if l == 0 or l == b else 1. 167 | # sum += eps_l_b / (1 - 4 * l ** 2) * np.cos(j * l * np.pi / b) 168 | # W[k, j] *= sum 169 | w = _clenshaw_curtis_weights(n=2 * b) 170 | W = np.empty((2 * b + 1, 2 * b + 2)) 171 | W[:] = w[:, None] 172 | elif grid_type == 'Gauss-Legendre': 173 | # We found this formula in: 174 | # "A Fast Algorithm for Spherical Grid Rotations and its Application to Singular Quadrature" 175 | # eq. 10 176 | _, w = leggauss(b + 1) 177 | W = w[:, None] * (2 * np.pi / (2 * b + 2) * np.ones(2 * b + 2)[None, :]) 178 | elif grid_type == 'SOFT': 179 | print("WARNING: SOFT quadrature weights don't work yet") 180 | k = np.arange(0, b) 181 | w = np.array([(2. / b) * np.sin(np.pi * (2. * j + 1.) / (4. * b)) * 182 | (np.sum((1. / (2 * k + 1)) 183 | * np.sin((2 * j + 1) * (2 * k + 1) 184 | * np.pi / (4. * b)))) 185 | for j in range(2 * b)]) 186 | W = w[:, None] * np.ones(2 * b)[None, :] 187 | else: 188 | raise ValueError('Unknown grid_type:' + str(grid_type)) 189 | 190 | return W 191 | 192 | 193 | def integrate(f, normalize=True): 194 | """ 195 | Integrate a function f : S^2 -> R over the sphere S^2, using the invariant integration measure 196 | mu((beta, alpha)) = sin(beta) dbeta dalpha 197 | i.e. this returns 198 | int_S^2 f(x) dmu(x) = int_0^2pi int_0^pi f(beta, alpha) sin(beta) dbeta dalpha 199 | 200 | :param f: a function of two scalar variables returning a scalar. 201 | :return: the integral of f over the 2-sphere 202 | """ 203 | from scipy.integrate import quad 204 | 205 | f2 = lambda alpha: quad(lambda beta: f(beta, alpha) * np.sin(beta), 206 | a=0, 207 | b=np.pi)[0] 208 | integral = quad(f2, 0, 2 * np.pi)[0] 209 | 210 | if normalize: 211 | return integral / (4 * np.pi) 212 | else: 213 | return integral 214 | 215 | 216 | def integrate_quad(f, grid_type, normalize=True, w=None): 217 | """ 218 | Integrate a function f : S^2 -> R, sampled on a grid of type grid_type, using quadrature weights w. 219 | 220 | :param f: an ndarray containing function values on a grid 221 | :param grid_type: the type of grid used to sample f 222 | :param normalize: whether to use the normalized Haar measure or not 223 | :param w: the quadrature weights. If not given, they are computed. 224 | :return: the integral of f over S^2. 225 | """ 226 | 227 | if grid_type != 'Gauss-Legendre' and grid_type != 'Clenshaw-Curtis': 228 | raise NotImplementedError 229 | 230 | b = (f.shape[1] - 2) // 2 # This works for Gauss-Legendre and Clenshaw-Curtis 231 | 232 | if w is None: 233 | w = quadrature_weights(b, grid_type) 234 | 235 | integral = np.sum(f * w) 236 | 237 | if normalize: 238 | return integral / (4 * np.pi) 239 | else: 240 | return integral 241 | 242 | 243 | def plot_sphere_func(f, grid='Clenshaw-Curtis', beta=None, alpha=None, colormap='jet', fignum=0, normalize=True): 244 | 245 | #TODO: All grids except Clenshaw-Curtis have holes at the poles 246 | # TODO: update this function now that we changed the order of axes in f 247 | 248 | import matplotlib 249 | matplotlib.use('WxAgg') 250 | matplotlib.interactive(True) 251 | from mayavi import mlab 252 | 253 | if normalize: 254 | f = (f - np.min(f)) / (np.max(f) - np.min(f)) 255 | 256 | if grid == 'Driscoll-Healy': 257 | b = f.shape[0] / 2 258 | elif grid == 'Clenshaw-Curtis': 259 | b = (f.shape[0] - 2) / 2 260 | elif grid == 'SOFT': 261 | b = f.shape[0] / 2 262 | elif grid == 'Gauss-Legendre': 263 | b = (f.shape[0] - 2) / 2 264 | 265 | if beta is None or alpha is None: 266 | beta, alpha = meshgrid(b=b, grid_type=grid) 267 | 268 | alpha = np.r_[alpha, alpha[0, :][None, :]] 269 | beta = np.r_[beta, beta[0, :][None, :]] 270 | f = np.r_[f, f[0, :][None, :]] 271 | 272 | x = np.sin(beta) * np.cos(alpha) 273 | y = np.sin(beta) * np.sin(alpha) 274 | z = np.cos(beta) 275 | 276 | mlab.figure(fignum, bgcolor=(1, 1, 1), fgcolor=(0, 0, 0), size=(600, 400)) 277 | mlab.clf() 278 | mlab.mesh(x, y, z, scalars=f, colormap=colormap) 279 | 280 | #mlab.view(90, 70, 6.2, (-1.3, -2.9, 0.25)) 281 | mlab.show() 282 | 283 | 284 | def plot_sphere_func2(f, grid='Clenshaw-Curtis', beta=None, alpha=None, colormap='jet', fignum=0, normalize=True): 285 | # TODO: update this function now that we have changed the order of axes in f 286 | import matplotlib.pyplot as plt 287 | from matplotlib import cm, colors 288 | from mpl_toolkits.mplot3d import Axes3D 289 | import numpy as np 290 | from scipy.special import sph_harm 291 | 292 | if normalize: 293 | f = (f - np.min(f)) / (np.max(f) - np.min(f)) 294 | 295 | if grid == 'Driscoll-Healy': 296 | b = f.shape[0] // 2 297 | elif grid == 'Clenshaw-Curtis': 298 | b = (f.shape[0] - 2) // 2 299 | elif grid == 'SOFT': 300 | b = f.shape[0] // 2 301 | elif grid == 'Gauss-Legendre': 302 | b = (f.shape[0] - 2) // 2 303 | 304 | if beta is None or alpha is None: 305 | beta, alpha = meshgrid(b=b, grid_type=grid) 306 | 307 | alpha = np.r_[alpha, alpha[0, :][None, :]] 308 | beta = np.r_[beta, beta[0, :][None, :]] 309 | f = np.r_[f, f[0, :][None, :]] 310 | 311 | x = np.sin(beta) * np.cos(alpha) 312 | y = np.sin(beta) * np.sin(alpha) 313 | z = np.cos(beta) 314 | 315 | # m, l = 2, 3 316 | # Calculate the spherical harmonic Y(l,m) and normalize to [0,1] 317 | # fcolors = sph_harm(m, l, beta, alpha).real 318 | # fmax, fmin = fcolors.max(), fcolors.min() 319 | # fcolors = (fcolors - fmin) / (fmax - fmin) 320 | print(x.shape, f.shape) 321 | 322 | if f.ndim == 2: 323 | f = cm.gray(f) 324 | print('2') 325 | 326 | # Set the aspect ratio to 1 so our sphere looks spherical 327 | fig = plt.figure(figsize=plt.figaspect(1.)) 328 | ax = fig.add_subplot(111, projection='3d') 329 | ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=f ) # cm.gray(f)) 330 | # Turn off the axis planes 331 | ax.set_axis_off() 332 | plt.show() 333 | 334 | 335 | def _clenshaw_curtis_weights(n): 336 | """ 337 | Computes the Clenshaw-Curtis quadrature using a fast FFT method. 338 | 339 | This is a 'brainless' port of MATLAB code found in: 340 | Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules 341 | Jorg Waldvogel, 2005 342 | http://www.sam.math.ethz.ch/~joergw/Papers/fejer.pdf 343 | 344 | :param n: 345 | :return: 346 | """ 347 | from scipy.fftpack import ifft, fft, fftshift 348 | 349 | # TODO python3 handles division differently from python2. Check how MATLAB interprets /, and if this code is still correct for python3 350 | 351 | # function [wf1,wf2,wcc] = fejer(n) 352 | # Weights of the Fejer2, Clenshaw-Curtis and Fejer1 quadratures by DFTs 353 | # n>1. Nodes: x_k = cos(k*pi/n) 354 | # N = [1:2:n-1]'; l=length(N); m=n-l; K=[0:m-1]'; 355 | N = np.arange(start=1, stop=n, step=2)[:, None] 356 | l = N.size 357 | m = n - l 358 | K = np.arange(start=0, stop=m)[:, None] 359 | 360 | # Fejer2 nodes: k=0,1,...,n; weights: wf2, wf2_n=wf2_0=0 361 | # v0 = [2./N./(N-2); 1/N(end); zeros(m,1)]; 362 | v0 = np.vstack([2. / N / (N-2), 1. / N[-1]] + [0] * m) 363 | 364 | # v2 = -v0(1:end-1) - v0(end:-1:2); 365 | # wf2 = ifft(v2); 366 | v2 = -v0[:-1] - v0[:0:-1] 367 | 368 | # Clenshaw-Curtis nodes: k=0,1,...,n; weights: wcc, wcc_n=wcc_0 369 | # g0 = -ones(n,1); 370 | g0 = -np.ones((n, 1)) 371 | 372 | # g0(1 + l) = g0(1 + l) + n; 373 | g0[l] = g0[l] + n 374 | 375 | # g0(1+m) = g0(1 + m) + n; 376 | g0[m] = g0[m] + n 377 | 378 | # g = g0/(n^2-1+mod(n,2)); 379 | g = g0 / (n ** 2 - 1 + n % 2) 380 | 381 | # wcc=ifft(v2 + g); 382 | wcc = ifft((v2 + g).flatten()).real 383 | wcc = np.hstack([wcc, wcc[0]]) 384 | 385 | # Fejer1 nodes: k=1/2,3/2,...,n-1/2; vector of weights: wf1 386 | # v0=[2*exp(i*pi*K/n)./(1-4*K.^2); zeros(l+1,1)]; 387 | # v1=v0(1:end-1)+conj(v0(end:-1:2)); wf1=ifft(v1); 388 | # don't need these 389 | 390 | return wcc * np.pi / (n / 2 + 1) # adjust for different scaling of python vs MATLAB fft 391 | -------------------------------------------------------------------------------- /lie_learn/spaces/S3.py: -------------------------------------------------------------------------------- 1 | from functools import lru_cache 2 | 3 | import numpy as np 4 | import lie_learn.spaces.S2 as S2 5 | 6 | 7 | def change_coordinates(coords, p_from='C', p_to='S'): 8 | """ 9 | Change Spherical to Cartesian coordinates and vice versa, for points x in S^3. 10 | 11 | We use the following coordinate system: 12 | https://en.wikipedia.org/wiki/N-sphere#Spherical_coordinates 13 | Except that we use the order (alpha, beta, gamma), where beta ranges from 0 to pi while alpha and gamma range from 14 | 0 to 2 pi. 15 | 16 | x0 = r * cos(alpha) 17 | x1 = r * sin(alpha) * cos(gamma) 18 | x2 = r * sin(alpha) * sin(gamma) * cos(beta) 19 | x3 = r * sin(alpha * sin(gamma) * sin(beta) 20 | 21 | :param conversion: 22 | :param coords: 23 | :return: 24 | """ 25 | if p_from == p_to: 26 | return coords 27 | elif p_from == 'S' and p_to == 'C': 28 | 29 | alpha = coords[..., 0] 30 | beta = coords[..., 1] 31 | gamma = coords[..., 2] 32 | r = 1. 33 | 34 | out = np.empty(alpha.shape + (4,)) 35 | 36 | ca = np.cos(alpha) 37 | cb = np.cos(beta) 38 | cc = np.cos(gamma) 39 | sa = np.sin(alpha) 40 | sb = np.sin(beta) 41 | sc = np.sin(gamma) 42 | out[..., 0] = r * ca 43 | out[..., 1] = r * sa * cc 44 | out[..., 2] = r * sa * sc * cb 45 | out[..., 3] = r * sa * sc * sb 46 | return out 47 | 48 | elif p_from == 'C' and p_to == 'S': 49 | 50 | raise NotImplementedError 51 | x = coords[..., 0] 52 | y = coords[..., 1] 53 | z = coords[..., 2] 54 | w = coords[..., 3] 55 | r = np.sqrt((coords ** 2).sum(axis=-1)) 56 | 57 | out = np.empty(x.shape + (3,)) 58 | out[..., 0] = np.arccos(z) # alpha 59 | out[..., 1] = np.arctan2(y, x) # beta 60 | out[..., 2] = np.arctan2(y, x) # gamma 61 | return out 62 | 63 | else: 64 | raise ValueError('Unknown conversion:' + str(p_from) + ' to ' + str(p_to)) 65 | 66 | 67 | def linspace(b, grid_type='SOFT'): 68 | """ 69 | Compute a linspace on the 3-sphere. 70 | 71 | Since S3 is ismorphic to SO(3), we use the grid grid_type from: 72 | FFTs on the Rotation Group 73 | Peter J. Kostelec and Daniel N. Rockmore 74 | http://www.cs.dartmouth.edu/~geelong/soft/03-11-060.pdf 75 | :param b: 76 | :return: 77 | """ 78 | # alpha = 2 * np.pi * np.arange(2 * b) / (2. * b) 79 | # beta = np.pi * (2 * np.arange(2 * b) + 1) / (4. * b) 80 | # gamma = 2 * np.pi * np.arange(2 * b) / (2. * b) 81 | 82 | beta, alpha = S2.linspace(b, grid_type) 83 | 84 | # According to this paper: 85 | # "Sampling sets and quadrature formulae on the rotation group" 86 | # We can just tack a sampling grid for S^1 to a sampling grid for S^2 to get a sampling grid for SO(3). 87 | gamma = 2 * np.pi * np.arange(2 * b) / (2. * b) 88 | 89 | return alpha, beta, gamma 90 | 91 | 92 | def meshgrid(b, grid_type='SOFT'): 93 | return np.meshgrid(*linspace(b, grid_type), indexing='ij') 94 | 95 | 96 | def integrate(f, normalize=True): 97 | """ 98 | Integrate a function f : S^3 -> R over the 3-sphere S^3, using the invariant integration measure 99 | mu((alpha, beta, gamma)) = dalpha sin(beta) dbeta dgamma 100 | i.e. this returns 101 | int_S^3 f(x) dmu(x) = int_0^2pi int_0^pi int_0^2pi f(alpha, beta, gamma) dalpha sin(beta) dbeta dgamma 102 | 103 | :param f: a function of three scalar variables returning a scalar. 104 | :param normalize: if we use the measure dalpha sin(beta) dbeta dgamma, 105 | the integral of f(a,b,c)=1 over the 3-sphere gives 8 pi^2. 106 | If normalize=True, we divide the result of integration by this normalization constant, so that f integrates to 1. 107 | In other words, use the normalized Haar measure. 108 | :return: the integral of f over the 3-sphere 109 | """ 110 | from scipy.integrate import quad 111 | 112 | f2 = lambda alpha, gamma: quad(lambda beta: f(alpha, beta, gamma) * np.sin(beta), 113 | a=0, 114 | b=np.pi)[0] 115 | f3 = lambda alpha: quad(lambda gamma: f2(alpha, gamma), 116 | a=0, 117 | b=2 * np.pi)[0] 118 | 119 | integral = quad(f3, 0, 2 * np.pi)[0] 120 | 121 | if normalize: 122 | return integral / (8 * np.pi ** 2) 123 | else: 124 | return integral 125 | 126 | 127 | def integrate_quad(f, grid_type, normalize=True, w=None): 128 | """ 129 | Integrate a function f : SO(3) -> R, sampled on a grid of type grid_type, using quadrature weights w. 130 | 131 | :param f: an ndarray containing function values on a grid 132 | :param grid_type: the type of grid used to sample f 133 | :param normalize: whether to use the normalized Haar measure or not 134 | :param w: the quadrature weights. If not given, they are computed. 135 | :return: the integral of f over S^2. 136 | """ 137 | 138 | if grid_type == 'SOFT': 139 | b = f.shape[0] // 2 140 | 141 | if w is None: 142 | w = quadrature_weights(b, grid_type) 143 | 144 | integral = np.sum(f * w[None, :, None]) 145 | else: 146 | raise NotImplementedError('Unsupported grid_type:', grid_type) 147 | 148 | if normalize: 149 | return integral 150 | else: 151 | return integral * 8 * np.pi ** 2 152 | 153 | 154 | @lru_cache(maxsize=32) 155 | def quadrature_weights(b, grid_type='SOFT'): 156 | """ 157 | Compute quadrature weights for the grid used by Kostelec & Rockmore [1, 2]. 158 | 159 | This grid is: 160 | alpha = 2 pi i / 2b 161 | beta = pi (2 j + 1) / 4b 162 | gamma = 2 pi k / 2b 163 | where 0 <= i, j, k < 2b are indices 164 | This grid can be obtained from the function: S3.linspace or S3.meshgrid 165 | 166 | The quadrature weights for this grid are 167 | w_B(j) = 2/b * sin(pi(2j + 1) / 4b) * sum_{k=0}^{b-1} 1 / (2 k + 1) sin((2j + 1)(2k + 1) pi / 4b) 168 | This is eq. 23 in [1] and eq. 2.15 in [2]. 169 | 170 | [1] SOFT: SO(3) Fourier Transforms 171 | Peter J. Kostelec and Daniel N. Rockmore 172 | 173 | [2] FFTs on the Rotation Group 174 | Peter J. Kostelec · Daniel N. Rockmore 175 | 176 | :param b: bandwidth (grid has shape 2b * 2b * 2b) 177 | :return: w: an array of length 2b containing the quadrature weigths 178 | """ 179 | if grid_type == 'SOFT': 180 | k = np.arange(0, b) 181 | w = np.array([(2. / b) * np.sin(np.pi * (2. * j + 1.) / (4. * b)) * 182 | (np.sum((1. / (2 * k + 1)) 183 | * np.sin((2 * j + 1) * (2 * k + 1) 184 | * np.pi / (4. * b)))) 185 | for j in range(2 * b)]) 186 | 187 | # This is not in the SOFT documentation, but we found that it is necessary to divide by this factor to 188 | # get correct results. 189 | w /= 2. * ((2 * b) ** 2) 190 | 191 | # In the SOFT source, they talk about the following weights being used for 192 | # odd-order transforms. Do not understand this, and the weights used above 193 | # (defined in the SOFT papers) seems to work. 194 | # w = np.array([(2. / b) * 195 | # (np.sum((1. / (2 * k + 1)) 196 | # * np.sin((2 * j + 1) * (2 * k + 1) 197 | # * np.pi / (4. * b)))) 198 | # for j in range(2 * b)]) 199 | return w 200 | else: 201 | raise NotImplementedError -------------------------------------------------------------------------------- /lie_learn/spaces/Tn.py: -------------------------------------------------------------------------------- 1 | """ 2 | The n-Torus 3 | """ 4 | 5 | import numpy as np 6 | 7 | def linspace(b, n=1, convention='regular'): 8 | if convention == 'regular': 9 | res = [] 10 | for i in range(n): 11 | res.append(np.arange(b) * 2 * np.pi / b) 12 | 13 | else: 14 | raise ValueError('Unknown convention:' + convention) 15 | 16 | return res -------------------------------------------------------------------------------- /lie_learn/spaces/__init__.py: -------------------------------------------------------------------------------- 1 | __author__ = 'tsc' 2 | -------------------------------------------------------------------------------- /lie_learn/spaces/rn.py: -------------------------------------------------------------------------------- 1 | """ 2 | n-dimensional real space, R^n. 3 | """ 4 | 5 | 6 | import numpy as np 7 | 8 | # The following functions are part of the public interface of this module; 9 | # other spaces / groups define their own meshgrid and linspace functions that work in an analogous way; 10 | # for R^n the standard numpy functions fulfill this role. 11 | from numpy import meshgrid, linspace 12 | 13 | 14 | def change_coordinates(coords, n, p_from='C', p_to='S'): 15 | """ 16 | Change Spherical to Cartesian coordinates and vice versa. 17 | 18 | todo: make this work for R^n and not just R^2, R^3 19 | 20 | :param conversion: 21 | :param coords: 22 | :return: 23 | """ 24 | 25 | coords = np.asarray(coords) 26 | 27 | if p_from == p_to: 28 | return coords 29 | 30 | if n == 2: 31 | if (p_from == 'P' or p_from == 'polar') and (p_to == 'C' or p_to == 'cartesian'): 32 | r = coords[..., 0] 33 | theta = coords[..., 1] 34 | out = np.empty_like(coords) 35 | out[..., 0] = r * np.cos(theta) 36 | out[..., 1] = r * np.sin(theta) 37 | return out 38 | elif (p_from == 'C' or p_from == 'cartesian') and (p_to == 'P' or p_to == 'polar'): 39 | x = coords[..., 0] 40 | y = coords[..., 1] 41 | out = np.empty_like(coords) 42 | out[..., 0] = np.sqrt(x ** 2 + y ** 2) 43 | out[..., 1] = np.arctan2(y, x) 44 | return out 45 | elif (p_from == 'C' or p_from == 'cartesian') and (p_to == 'H' or p_to == 'homogeneous'): 46 | x = coords[..., 0] 47 | y = coords[..., 1] 48 | out = np.empty(coords.shape[:-1] + (3,)) 49 | out[..., 0] = x 50 | out[..., 1] = y 51 | out[..., 2] = 1. 52 | return out 53 | elif (p_from == 'H' or p_from == 'homogeneous') and (p_to == 'C' or p_to == 'cartesian'): 54 | xc = coords[..., 0] 55 | yc = coords[..., 1] 56 | c = coords[..., 2] 57 | out = np.empty(coords.shape[:-1] + (2,)) 58 | out[..., 0] = xc / c 59 | out[..., 1] = yc / c 60 | return out 61 | else: 62 | raise ValueError('Unknown conversion' + str(p_from) + ' to ' + str(p_to)) 63 | 64 | elif n == 3: 65 | 66 | if p_from == 'S' and p_to == 'C': 67 | 68 | theta = coords[..., 0] 69 | phi = coords[..., 1] 70 | r = coords[..., 2] 71 | 72 | out = np.empty(theta.shape + (3,)) 73 | 74 | ct = np.cos(theta) 75 | cp = np.cos(phi) 76 | st = np.sin(theta) 77 | sp = np.sin(phi) 78 | out[..., 0] = r * st * cp # x 79 | out[..., 1] = r * st * sp # y 80 | out[..., 2] = r * ct # z 81 | return out 82 | 83 | elif p_from == 'C' and p_to == 'S': 84 | 85 | x = coords[..., 0] 86 | y = coords[..., 1] 87 | z = coords[..., 2] 88 | 89 | out = np.empty_like(coords) 90 | out[..., 2] = np.sqrt(x ** 2 + y ** 2 + z ** 2) # r 91 | out[..., 0] = np.arccos(z / out[..., 2]) # theta 92 | out[..., 1] = np.arctan2(y, x) # phi 93 | return out 94 | 95 | else: 96 | raise ValueError('Unknown conversion:' + str(p_from) + ' to ' + str(p_to)) 97 | else: 98 | raise ValueError('Only dimension n=2 and n=3 supported for now.') 99 | 100 | 101 | def linspace(b, convention): 102 | 103 | pass 104 | 105 | 106 | 107 | -------------------------------------------------------------------------------- /lie_learn/spaces/spherical_quadrature.pyx: -------------------------------------------------------------------------------- 1 | 2 | from lie_learn.representations.SO3.spherical_harmonics import rsh 3 | 4 | import numpy as np 5 | cimport numpy as np 6 | 7 | 8 | def estimate_spherical_quadrature_weights(sampling_set, max_bandwidth, 9 | normalization='quantum', condon_shortley=True, 10 | verbose=True): 11 | """ 12 | 13 | :param sampling_set: 14 | :param max_bandwith: 15 | :return: 16 | """ 17 | cdef int l 18 | cdef int m 19 | cdef int ll 20 | cdef int mm 21 | cdef int i 22 | 23 | cdef int M = sampling_set.shape[0] 24 | cdef int N = max_bandwidth 25 | cdef int N_total = (N + 1) ** 2 # = sum_l=0^N (2l + 1) 26 | 27 | cdef np.ndarray[np.float64_t, ndim=2] l_array = np.empty((N_total, 1)) 28 | cdef np.ndarray[np.float64_t, ndim=2] m_array = np.empty((N_total, 1)) 29 | 30 | theta = sampling_set[:, 0] 31 | phi = sampling_set[:, 1] 32 | 33 | if verbose: 34 | print 'Computing index arrays...' 35 | 36 | i = 0 37 | for l in range(N + 1): 38 | for m in range(-l, l + 1): 39 | l_array[i, 0] = l 40 | m_array[i, 0] = m 41 | i += 1 42 | 43 | if verbose: 44 | print 'Computing spherical harmonics...' 45 | Y = rsh(l_array, m_array, theta[None, :], phi[None, :], 46 | normalization=normalization, condon_shortley=condon_shortley) 47 | 48 | if verbose: 49 | print 'Computing least squares input' 50 | B = np.empty((N_total ** 2, M)) 51 | t = np.empty(N_total ** 2) 52 | i = 0 53 | 54 | #print M, N, N_total 55 | #print theta[None, :].shape 56 | #print phi[None, :].shape 57 | #print Y.shape 58 | #print B.shape 59 | #print t.shape 60 | 61 | for l in range(N + 1): 62 | for m in range(-l, l + 1): 63 | rlm = Y[l ** 2 + l + m, :] 64 | for ll in range(N + 1): 65 | for mm in range(-ll, ll + 1): 66 | B[i, :] = rlm * Y[ll ** 2 + ll + mm, :] 67 | t[i] = float(ll == l and mm == m) 68 | i += 1 69 | 70 | if verbose: 71 | print 'Computing least squares solution' 72 | return np.linalg.lstsq(B, t) 73 | -------------------------------------------------------------------------------- /lie_learn/spectral/FFTBase.py: -------------------------------------------------------------------------------- 1 | 2 | class FFTBase(object): 3 | 4 | def __init__(self): 5 | pass 6 | 7 | def analyze(self, f): 8 | raise NotImplementedError('FFTBase.analyze should be implemented in subclass') 9 | 10 | def synthesize(self, f_hat): 11 | raise NotImplementedError('FFTBase.synthesize should be implemented in subclass') 12 | -------------------------------------------------------------------------------- /lie_learn/spectral/PolarFFT.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from .FFTBase import FFTBase 4 | from pynfft.nfft import NFFT 5 | 6 | 7 | # UNFINISHED 8 | 9 | class PolarFFT(FFTBase): 10 | 11 | def __init__(self, nx, ny, nt, nr): 12 | 13 | # Initialize the non-equispaced FFT 14 | self.nfft = NFFT(N=(nx, ny), M=nx * ny, n=None, m=12, flags=None) 15 | 16 | # Set up the polar sampling grid 17 | theta = np.linspace(0, 2 * np.pi, nt) 18 | r = np.linspace(0, 1., nr) 19 | T, R = np.meshgrid(theta, r) 20 | self.nfft.x = np.c_[T[..., None], R[..., None]].flatten() 21 | self.nfft.precompute() 22 | 23 | def analyze(self, f): 24 | 25 | self.nfft.f_hat = f 26 | f_hat = self.nfft.forward() 27 | 28 | return f_hat 29 | 30 | def synthesize(self, f_hat): 31 | 32 | self.nfft.f = f_hat 33 | f = self.nfft.adjoint() 34 | 35 | return f_hat 36 | -------------------------------------------------------------------------------- /lie_learn/spectral/S2FFT.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from scipy.fftpack import fft, ifft, fftshift 4 | 5 | from lie_learn.spectral.FFTBase import FFTBase 6 | import lie_learn.spaces.S2 as S2 7 | from lie_learn.representations.SO3.spherical_harmonics import csh, sh 8 | 9 | 10 | class S2_FT_Naive(FFTBase): 11 | """ 12 | The most naive implementation of the discrete spherical Fourier transform: 13 | explicitly construct the Fourier matrix F and multiply by it to perform the Fourier transform. 14 | """ 15 | 16 | def __init__(self, L_max, 17 | grid_type='Gauss-Legendre', 18 | field='real', normalization='quantum', condon_shortley='cs'): 19 | 20 | super().__init__() 21 | 22 | self.b = L_max + 1 23 | 24 | # Compute a grid of spatial sampling points and associated quadrature weights 25 | beta, alpha = S2.meshgrid(b=self.b, grid_type=grid_type) 26 | self.w = S2.quadrature_weights(b=self.b, grid_type=grid_type) 27 | self.spatial_grid_shape = beta.shape 28 | self.num_spatial_points = beta.size 29 | 30 | # Determine for which degree and order we want the spherical harmonics 31 | irreps = np.arange(self.b) # TODO find out upper limit for exact integration for each grid type 32 | ls = [[ls] * (2 * ls + 1) for ls in irreps] 33 | ls = np.array([ll for sublist in ls for ll in sublist]) # 0, 1, 1, 1, 2, 2, 2, 2, 2, ... 34 | ms = [list(range(-ls, ls + 1)) for ls in irreps] 35 | ms = np.array([mm for sublist in ms for mm in sublist]) # 0, -1, 0, 1, -2, -1, 0, 1, 2, ... 36 | self.num_spectral_points = ms.size # This equals sum_{l=0}^{b-1} 2l+1 = b^2 37 | 38 | # In one shot, sample the spherical harmonics at all spectral (l, m) and spatial (beta, alpha) coordinates 39 | self.Y = sh(ls[None, None, :], ms[None, None, :], beta[:, :, None], alpha[:, :, None], 40 | field=field, normalization=normalization, condon_shortley=condon_shortley == 'cs') 41 | 42 | # Convert to a matrix 43 | self.Ymat = self.Y.reshape(self.num_spatial_points, self.num_spectral_points) 44 | 45 | def analyze(self, f): 46 | return self.Ymat.T.conj().dot((f * self.w).flatten()) 47 | 48 | def synthesize(self, f_hat): 49 | return self.Ymat.dot(f_hat).reshape(self.spatial_grid_shape) 50 | 51 | 52 | def setup_legendre_transform(b): 53 | """ 54 | Compute a set of matrices containing coefficients to be used in a discrete Legendre transform. 55 | 56 | The discrete Legendre transform of a data vector s[k] (k=0, ..., 2b-1) is defined as 57 | 58 | s_hat(l, m) = sum_{k=0}^{2b-1} P_l^m(cos(beta_k)) s[k] 59 | for l = 0, ..., b-1 and -l <= m <= l, 60 | where P_l^m is the associated Legendre function of degree l and order m, 61 | beta_k = ... 62 | 63 | Computing Fourier Transforms and Convolutions on the 2-Sphere 64 | J.R. Driscoll, D.M. Healy 65 | 66 | FFTs for the 2-Sphere–Improvements and Variations 67 | D.M. Healy, Jr., D.N. Rockmore, P.J. Kostelec, and S. Moore 68 | 69 | :param b: bandwidth of the transform 70 | :return: lt, an array of shape (N, 2b), containing samples of the Legendre functions, 71 | where N is the number of spectral points for a signal of bandwidth b. 72 | """ 73 | dim = np.sum(np.arange(b) * 2 + 1) 74 | lt = np.empty((2 * b, dim)) 75 | 76 | beta, _ = S2.linspace(b, grid_type='Driscoll-Healy') 77 | sample_points = np.cos(beta) 78 | 79 | # TODO move quadrature weight computation to S2.py 80 | weights = [(1. / b) * np.sin(np.pi * j * 0.5 / b) * 81 | np.sum([1. / (2 * l + 1) * np.sin((2 * l + 1) * np.pi * j * 0.5 / b) 82 | for l in range(b)]) 83 | for j in range(2 * b)] 84 | weights = np.array(weights) 85 | 86 | zeros = np.zeros_like(sample_points) 87 | i = 0 88 | for l in range(b): 89 | for m in range(-l, l + 1): 90 | # Z = np.sqrt(((2 * l + 1) * factorial(l - m)) / float(4 * np.pi * factorial(l + m))) * np.pi / 2 91 | # lt[i, :] = lpmv(m, l, sample_points) * weights * Z 92 | 93 | # The spherical harmonics code appears to be more stable than the (unnormalized) associated Legendre 94 | # function code. 95 | # (Note: the spherical harmonics evaluated at alpha=0 is the associated Legendre function)) 96 | lt[:, i] = csh(l, m, beta, zeros, normalization='seismology').real * weights * np.pi / 2 97 | 98 | i += 1 99 | 100 | return lt 101 | 102 | 103 | def setup_legendre_transform_indices(b): 104 | ms = [list(range(-ls, ls + 1)) for ls in range(b)] 105 | ms = [mm for sublist in ms for mm in sublist] # 0, -1, 0, 1, -2, -1, 0, 1, 2, ... 106 | ms = [mm % (2 * b) for mm in ms] 107 | return ms 108 | 109 | 110 | def sphere_fft(f, lt=None, lti=None): 111 | """ 112 | Compute the Spherical Fourier transform of f. 113 | We use complex, seismology-normalized, centered spherical harmonics, which are orthonormal (see rep_bases.py). 114 | 115 | The spherical Fourier transform is defined: 116 | \hat{f}_l^m = int_0^pi dbeta sin(beta) int_0^2pi dalpha f(beta, alpha) Y_l^{m*}(beta, alpha) 117 | (where we use the invariant area element dOmega = sin(beta) dbeta dalpha for the 2-sphere) 118 | 119 | We have Y_l^m(beta, alpha) = P_l^m(cos(beta)) * e^{im alpha}, where P_l^m is the associated Legendre function, 120 | so we can rewrite: 121 | \hat{f}_l^m = int_0^pi dbeta sin(beta) (int_0^2pi dalpha f(beta, alpha) e^{im alpha} ) P_l^m(cos(beta)) 122 | 123 | The integral over alpha can be evaluated by FFT: 124 | \bar{f}(beta_k, m) = int_0^2pi dalpha f(beta_k, alpha) e^{im alpha} = FFT(f, axis=1)[beta_k, m] 125 | 126 | Then we have 127 | \hat{f}_l^m = int_0^pi sin(beta) dbeta \bar{f}(beta, m) P_l^m(cos(beta)) 128 | = sum_k \bar{f}[beta_k, m] P_l^m(cos(beta_k)) w_k 129 | For appropriate quadrature weights w_k. This sum is called the discrete Legendre transform of \bar{f} 130 | 131 | We return \hat{f} as a flat vector. Hence, the precomputed P_l^m(cos(beta_k)) w_k is stored as an array of with a 132 | combined (l, m)-axis and a k axis. We bring the data \bar{f}[beta_k, m] into the same form, by indexing with lti 133 | and then reduce over the beta_k axis. 134 | 135 | Main source: 136 | Engineering Applications of Noncommutative Harmonic Analysis. 137 | 4.7.2 - Orthogonal Expansions on the Sphere 138 | G.S. Chrikjian, A.B. Kyatkin 139 | 140 | Further information: 141 | SOFT: SO(3) Fourier Transforms 142 | Peter J. Kostelec and Daniel N. Rockmore 143 | 144 | Generalized FFTs-a survey of some recent results 145 | Maslen & Rockmore 146 | 147 | Computing Fourier transforms and convolutions on the 2-sphere. 148 | Driscoll, J., & Healy, D. (1994). 149 | 150 | :param f: array of samples of the function to be transformed. Shape (2 * b, 2 * b). grid_type: Driscoll-Healy 151 | :param lt: precomputed Legendre transform matrices, from setup_legendre_transform(). 152 | :param lti: precomputed Legendre transform indices, from setup_legendre_transform_indices(). 153 | :return: f_hat, the spherical Fourier transform of f. This is an array of size sum_l=0^{b-1} 2 l + 1. 154 | the coefficients are ordered as (l=0, m=0), (l=1, m=-1), (l=1, m=0), (l=1,m=1), ... 155 | """ 156 | assert f.shape[-2] == f.shape[-1] 157 | assert f.shape[-2] % 2 == 0 158 | b = f.shape[-2] // 2 159 | 160 | if lt is None: 161 | lt = setup_legendre_transform(b) 162 | 163 | if lti is None: 164 | lti = setup_legendre_transform_indices(b) 165 | 166 | # First, FFT along the alpha axis (last axis) 167 | # This gives the array f_bar with axes for beta and m. 168 | f_bar = fft(f, axis=-1) 169 | 170 | # Perform Legendre transform 171 | f_hat = (f_bar[..., lti] * lt).sum(axis=-2) 172 | return f_hat 173 | -------------------------------------------------------------------------------- /lie_learn/spectral/S2FFT_NFFT.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from .FFTBase import FFTBase 4 | from pynfft import nfsft 5 | from lie_learn.spaces.spherical_quadrature import estimate_spherical_quadrature_weights 6 | from lie_learn.representations.SO3.irrep_bases import change_of_basis_matrix, change_of_basis_function 7 | 8 | class S2FFT_NFFT(FFTBase): 9 | 10 | def __init__(self, L_max, x, w=None): 11 | """ 12 | 13 | :param L_max: maximum spherical harmonic degree 14 | :param x: coordinates on spherical / spatial grid 15 | :param w: quadrature weights for the grid x 16 | """ 17 | 18 | # If x is a list (generated by S2.meshgrid), convert to (M, 2) array 19 | if isinstance(x, list): 20 | x = np.c_[x[0].flatten()[:, None], x[1].flatten()[:, None]] 21 | 22 | # The NFSFT class can synthesis / analyze functions in terms of 23 | # NFFT-normalized, centered, complex spherical harmonics without Condon-Shortley phase. 24 | self._nfsft = nfsft.NFSFT(N=L_max, x=x) 25 | 26 | # Compute a change-of-basis matrix from the NFFT spherical harmonics to our prefered choice, the 27 | # quantum-normalized, centered, real spherical harmonics with Condon-Shortley phase. 28 | #TODO: change this to change_of_basis_function (test that it works..) 29 | #self._c2r = change_of_basis_matrix(np.arange(L_max + 1), 30 | # frm=('complex', 'nfft', 'centered', 'nocs'), 31 | # to=('real', 'quantum', 'centered', 'cs')) 32 | #self._r2c = change_of_basis_matrix(np.arange(L_max + 1), 33 | # to=('complex', 'nfft', 'centered', 'nocs'), 34 | # frm=('real', 'quantum', 'centered', 'cs')) 35 | #self._c = change_of_basis_matrix(np.arange(L_max + 1), 36 | # frm=('real', 'nfft', 'centered', 'cs'), 37 | # to=('complex', 'quantum', 'centered', 'nocs')) 38 | 39 | 40 | 41 | self._c2r_func = change_of_basis_function(np.arange(L_max + 1), 42 | frm=('complex', 'nfft', 'centered', 'nocs'), 43 | to=('real', 'quantum', 'centered', 'cs')) 44 | #self._r2c_func = change_of_basis_function(np.arange(L_max + 1), 45 | # frm=('real', 'quantum', 'centered', 'cs'), 46 | # to=('complex', 'nfft', 'centered', 'nocs')) 47 | 48 | # In the synthesize() function, we will need c2r.conj().T as a function (not a matrix). 49 | # It happens to be the case that the following is equal to c2r.conj().T: 50 | c2r_conj_T = change_of_basis_function(np.arange(L_max + 1), 51 | frm=('real', 'nfft', 'centered', 'cs'), 52 | to=('complex', 'quantum', 'centered', 'nocs')) 53 | self._c2r_T = lambda vec: c2r_conj_T(vec.conj()).conj() 54 | 55 | if w is None: 56 | # Precompute quadrature weights 57 | self.w = estimate_spherical_quadrature_weights( 58 | sampling_set=x, max_bandwidth=L_max, 59 | normalization='quantum', condon_shortley=True)[0] 60 | else: 61 | self.w = w.flatten() 62 | 63 | self.x = x 64 | self.L_max = L_max 65 | 66 | def analyze(self, f): 67 | 68 | # We want to perform the *weighted* adjoint FFT, so that we get the exact Fourier coefficients 69 | # (at least for a proper sampling grid such as Clenshaw-Curtis or Gauss-Legendre and the respective weights) 70 | # Hence, the function to be transformed is f * w 71 | self._nfsft.f = f * self.w 72 | 73 | # Expand the weighted function in terms of the conjugate of 74 | # NFFT-normalized, centered, complex spherical harmonics without Condon-Shortley phase: 75 | # a_lm = sum_i=0^M Y_lm(theta_i, phi_i).conj() * w_i * f(theta_i, phi_i) 76 | self._nfsft.adjoint() 77 | 78 | # The computed Fourier components a_lm are with respect to the basis of NFFT spherical harmonics, 79 | # so change the basis. 80 | # Let Y denote the M by (L_max+1)^2 matrix of NFFT spherical harmonics. 81 | # then a = Y.conj().T.dot(f * w), as computed by _nfsft.adjoint() 82 | # Since, Y.conj().T = r2c.conj().dot(R.T), we have a = r2c.conj().dot(R.T.dot(f * w)) 83 | # To cancel the r2c.conj(), we multiply with c2r.conj() 84 | #a = self._c2r.conj().dot(self._nfsft.get_f_hat_flat()).real 85 | #b = self._c2r_func(self._nfsft.get_f_hat_flat().conj()).conj().real 86 | #print 'DIFF', np.sum(np.abs(a-b)) 87 | #return self._c2r.conj().dot(self._nfsft.get_f_hat_flat()).real 88 | 89 | return self._c2r_func(self._nfsft.get_f_hat_flat().conj()).conj().real 90 | 91 | def synthesize(self, f_hat): 92 | # self._nfsft.trafo() computes the synthesis / forward transform using NFFT complex SH: 93 | # f = Y f_hat, where Y is the M by (L_max+1)^2 matrix of complex NFFT spherical harmonics. 94 | # We have R.T = c2r.dot(Y.T), so f = R.dot(f_hat) = Y.dot(c2r.T.dot(f_hat)) 95 | #cfh = self._c2r.T.dot(f_hat) 96 | cfh = self._c2r_T(f_hat) 97 | self._nfsft.set_f_hat_flat(cfh) 98 | f = self._nfsft.trafo(use_dft=False, return_copy=True) 99 | return f.real -------------------------------------------------------------------------------- /lie_learn/spectral/S2_conv.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | import lie_learn.spaces.S2 as S2 5 | import lie_learn.groups.SO3 as SO3 6 | from lie_learn.spectral.S2FFT import S2_FT_Naive 7 | from lie_learn.spectral.SO3FFT_Naive import SO3_FT_Naive 8 | 9 | 10 | def conv_test(): 11 | """ 12 | 13 | :return: 14 | """ 15 | from lie_learn.spectral.SO3FFT_Naive import SO3_FT_Naive 16 | 17 | b = 10 18 | f1 = np.ones((2 * b + 2, b + 1)) 19 | f2 = np.ones((2 * b + 2, b + 1)) 20 | 21 | s2_fft = S2_FT_Naive(L_max=b - 1, grid_type='Gauss-Legendre', 22 | field='real', normalization='quantum', condon_shortley='cs') 23 | 24 | so3_fft = SO3_FT_Naive(L_max=b - 1, 25 | field='real', normalization='quantum', order='centered', condon_shortley='cs') 26 | 27 | # Spherical Fourier transform 28 | f1_hat = s2_fft.analyze(f1) 29 | f2_hat = s2_fft.analyze(f2) 30 | 31 | # Perform block-wise outer product 32 | f12_hat = [] 33 | for l in range(b): 34 | f1_hat_l = f1_hat[l ** 2:l ** 2 + 2 * l + 1] 35 | f2_hat_l = f2_hat[l ** 2:l ** 2 + 2 * l + 1] 36 | 37 | f12_hat_l = f1_hat_l[:, None] * f2_hat_l[None, :].conj() 38 | f12_hat.append(f12_hat_l) 39 | 40 | # Inverse SO(3) Fourier transform 41 | f12 = so3_fft.synthesize(f12_hat) 42 | 43 | return f12 44 | 45 | 46 | def spectral_S2_conv(f1, f2, s2_fft=None, so3_fft=None): 47 | """ 48 | Compute the convolution of two functions on the 2-sphere. 49 | Let f1 : S^2 -> R and f2 : S^2 -> R, then the convolution is defined as 50 | f1 * f2(g) = int_{S^2} f1(x) f2(g^{-1} x) dx, 51 | where g in SO(3) and dx is the normalized Haar measure on S^2. 52 | 53 | The convolution is computed by a Fourier transform. 54 | It can be shown that the SO(3)-Fourier transform of the convolution f1 * f2 is equal to the outer product 55 | of the spherical Fourier transform of f1 and f2. 56 | Specifically, let f1_hat be the spherical FT of f1 and f2_hat the spherical FT of f2. 57 | These vectors are split into chunks of dimension 2l+1, for l=0, ..., b (the bandwidth) 58 | For each degree, we take the outer product to obtain a (2l+1) x (2l+1) matrix, which is the degree-l 59 | block of the FT of f1*f2. 60 | For more details, see our note on "Convolution on S^2 and SO(3)" 61 | 62 | :param f1: 63 | :param f2: 64 | :param s2_fft: 65 | :param so3_fft: 66 | :return: 67 | """ 68 | 69 | b = f1.shape[1] - 1 # TODO we assume a Gauss-Legendre grid for S^2 here 70 | 71 | if s2_fft is None: 72 | s2_fft = S2_FT_Naive(L_max=b - 1, grid_type='Gauss-Legendre', 73 | field='real', normalization='quantum', condon_shortley='cs') 74 | 75 | if so3_fft is None: 76 | so3_fft = SO3_FT_Naive(L_max=b - 1, 77 | field='real', normalization='quantum', order='centered', condon_shortley='cs') 78 | 79 | # Spherical Fourier transform 80 | f1_hat = s2_fft.analyze(f1) 81 | f2_hat = s2_fft.analyze(f2) 82 | 83 | # Perform block-wise outer product 84 | f12_hat = [] 85 | for l in range(b): 86 | f1_hat_l = f1_hat[l ** 2:l ** 2 + 2 * l + 1] 87 | f2_hat_l = f2_hat[l ** 2:l ** 2 + 2 * l + 1] 88 | 89 | f12_hat_l = f1_hat_l[:, None] * f2_hat_l[None, :].conj() 90 | f12_hat.append(f12_hat_l) 91 | 92 | # Inverse SO(3) Fourier transform 93 | return so3_fft.synthesize(f12_hat) 94 | 95 | 96 | def naive_S2_conv(f1, f2, alpha, beta, gamma, g_parameterization='EA323'): 97 | """ 98 | Compute int_S^2 f1(x) f2(g^{-1} x)* dx, 99 | where x = (theta, phi) is a point on the sphere S^2, 100 | and g = (alpha, beta, gamma) is a point in SO(3) in Euler angle parameterization 101 | 102 | :param f1, f2: functions to be convolved 103 | :param alpha, beta, gamma: the rotation at which to evaluate the result of convolution 104 | :return: 105 | """ 106 | # This fails 107 | def integrand(theta, phi): 108 | g_inv = SO3.invert((alpha, beta, gamma), parameterization=g_parameterization) 109 | g_inv_theta, g_inv_phi, _ = SO3.transform_r3(g=g_inv, x=(theta, phi, 1.), 110 | g_parameterization=g_parameterization, x_parameterization='S') 111 | return f1(theta, phi) * f2(g_inv_theta, g_inv_phi).conj() 112 | 113 | return S2.integrate(f=integrand, normalize=True) 114 | 115 | 116 | def naive_S2_conv_v2(f1, f2, alpha, beta, gamma, g_parameterization='EA323'): 117 | """ 118 | Compute int_S^2 f1(x) f2(g^{-1} x)* dx, 119 | where x = (theta, phi) is a point on the sphere S^2, 120 | and g = (alpha, beta, gamma) is a point in SO(3) in Euler angle parameterization 121 | 122 | :param f1, f2: functions to be convolved 123 | :param alpha, beta, gamma: the rotation at which to evaluate the result of convolution 124 | :return: 125 | """ 126 | 127 | theta, phi = S2.meshgrid(b=3, grid_type='Gauss-Legendre') 128 | w = S2.quadrature_weights(b=3, grid_type='Gauss-Legendre') 129 | 130 | print(theta.shape, phi.shape) 131 | s2_coords = np.c_[theta[..., None], phi[..., None]] 132 | print(s2_coords.shape) 133 | r3_coords = np.c_[theta[..., None], phi[..., None], np.ones_like(theta)[..., None]] 134 | 135 | # g_inv = SO3.invert((alpha, beta, gamma), parameterization=g_parameterization) 136 | # g_inv = (-gamma, -beta, -alpha) 137 | g_inv = (alpha, beta, gamma) # wrong 138 | 139 | ginvx = SO3.transform_r3(g=g_inv, x=r3_coords, g_parameterization=g_parameterization, x_parameterization='S') 140 | print(ginvx.shape) 141 | g_inv_theta = ginvx[..., 0] 142 | g_inv_phi = ginvx[..., 1] 143 | g_inv_r = ginvx[..., 2] 144 | 145 | print(g_inv_theta, g_inv_phi, g_inv_r) 146 | 147 | f1_grid = f1(theta, phi) 148 | f2_grid = f2(g_inv_theta, g_inv_phi) 149 | 150 | print(f1_grid.shape, f2_grid.shape, w.shape) 151 | return np.sum(f1_grid * f2_grid * w) 152 | 153 | 154 | 155 | -------------------------------------------------------------------------------- /lie_learn/spectral/SO3_conv.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | import numpy as np 5 | 6 | from lie_learn.spectral.SO3FFT_Naive import SO3_FT_Naive 7 | 8 | 9 | def conv_test(): 10 | """ 11 | Compute the convolution of two functions on SO(3). 12 | Let f1 : SO(3) -> R and f2 : SO(3) -> R, then the convolution is defined as 13 | f1 * f2(g) = int_{SO(3)} f1(h) f2(g^{-1} h) dh, 14 | where g in SO(3) and dh is the normalized Haar measure on SO(3). 15 | 16 | The convolution is computed by a Fourier transform. 17 | It can be shown that the SO(3) Fourier transform of the convolution f1 * f2 is equal to the matrix product 18 | of the SO(3) Fourier transforms of f1 and f2. 19 | For more details, see the note on "Convolution on S^2 and SO(3)" 20 | 21 | :return: 22 | """ 23 | from lie_learn.spectral.SO3FFT_Naive import SO3_FT_Naive 24 | 25 | b = 10 26 | f1 = np.ones((2 * b + 2, b + 1)) #TODO 27 | f2 = np.ones((2 * b + 2, b + 1)) 28 | 29 | s2_fft = S2_FT_Naive(L_max=b - 1, grid_type='Gauss-Legendre', 30 | field='real', normalization='quantum', condon_shortley='cs') 31 | 32 | so3_fft = SO3_FT_Naive(L_max=b - 1, 33 | field='real', normalization='quantum', order='centered', condon_shortley='cs') 34 | 35 | # Spherical Fourier transform 36 | f1_hat = s2_fft.analyze(f1) 37 | f2_hat = s2_fft.analyze(f2) 38 | 39 | # Perform block-wise outer product 40 | f12_hat = [] 41 | for l in range(b): 42 | f1_hat_l = f1_hat[l ** 2:l ** 2 + 2 * l + 1] 43 | f2_hat_l = f2_hat[l ** 2:l ** 2 + 2 * l + 1] 44 | 45 | f12_hat_l = f1_hat_l[:, None] * f2_hat_l[None, :].conj() 46 | f12_hat.append(f12_hat_l) 47 | 48 | # Inverse SO(3) Fourier transform 49 | f12 = so3_fft.synthesize(f12_hat) 50 | 51 | return f12 52 | 53 | 54 | def SO3_convolve(f, g, dw=None, d=None): 55 | 56 | assert f.shape == g.shape 57 | assert f.shape[0] % 2 == 0 58 | b = f.shape[0] / 2 59 | 60 | if d is None: 61 | d = setup_d_transform(b) 62 | 63 | # To convolve, first perform a Fourier transform on f and g: 64 | F = SO3_fft(f, dw) 65 | G = SO3_fft(g, dw) 66 | 67 | # The Fourier transform of the convolution f*g is the matrix product FG 68 | # of their Fourier transforms F and G: 69 | FG = [np.dot(a, b) for (a, b) in zip(F, G)] 70 | 71 | # The convolution is obtain by inverse Fourier transforming FG: 72 | return SO3_ifft(FG, d) -------------------------------------------------------------------------------- /lie_learn/spectral/T1FFT.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from numpy.fft import fft, ifft, fftshift, ifftshift 4 | from .FFTBase import FFTBase 5 | 6 | 7 | class T1FFT(FFTBase): 8 | """ 9 | The Fast Fourier Transform on the Circle / 1-Torus / 1-Sphere. 10 | """ 11 | 12 | @staticmethod 13 | def analyze(f, axis=0): 14 | """ 15 | Compute the Fourier Transform of the discretely sampled function f : T^1 -> C. 16 | 17 | Let f : T^1 -> C be a band-limited function on the circle. 18 | The samples f(theta_k) correspond to points on a regular grid on the circle, as returned by spaces.T1.linspace: 19 | theta_k = 2 pi k / N 20 | for k = 0, ..., N - 1 21 | 22 | This function computes 23 | \hat{f}_n = (1/N) \sum_{k=0}^{N-1} f(theta_k) e^{-i n theta_k} 24 | which, if f has band-limit less than N, is equal to: 25 | \hat{f}_n = \int_0^{2pi} f(theta) e^{-i n theta} dtheta / 2pi, 26 | = 27 | where dtheta / 2pi is the normalized Haar measure on T^1, and < , > denotes the inner product on Hilbert space, 28 | with respect to which this transform is unitary. 29 | 30 | The range of frequencies n is -floor(N/2) <= n <= ceil(N/2) - 1 31 | 32 | :param f: 33 | :param axis: 34 | :return: 35 | """ 36 | # The numpy FFT returns coefficients in a different order than we want them, 37 | # and using a different normalization. 38 | fhat = fft(f, axis=axis) 39 | fhat = fftshift(fhat, axes=axis) 40 | return fhat / f.shape[axis] 41 | 42 | @staticmethod 43 | def synthesize(f_hat, axis=0): 44 | """ 45 | Compute the inverse / synthesis Fourier transform of the function f_hat : Z -> C. 46 | The function f_hat(n) is sampled at points in a limited range -floor(N/2) <= n <= ceil(N/2) - 1 47 | 48 | This function returns 49 | f[k] = f(theta_k) = sum_{n=-floor(N/2)}^{ceil(N/2)-1} f_hat(n) exp(i n theta_k) 50 | where theta_k = 2 pi k / N 51 | for k = 0, ..., N - 1 52 | 53 | :param f_hat: 54 | :param axis: 55 | :return: 56 | """ 57 | 58 | f_hat = ifftshift(f_hat * f_hat.shape[axis], axes=axis) 59 | f = ifft(f_hat, axis=axis) 60 | return f 61 | 62 | @staticmethod 63 | def analyze_naive(f): 64 | f_hat = np.zeros_like(f) 65 | for n in range(f.size): 66 | for k in range(f.size): 67 | theta_k = k * 2 * np.pi / f.size 68 | f_hat[n] += f[k] * np.exp(-1j * n * theta_k) 69 | return fftshift(f_hat / f.size, axes=0) 70 | -------------------------------------------------------------------------------- /lie_learn/spectral/T2FFT.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from numpy.fft import fft2, ifft2, fftshift, ifftshift 4 | from .FFTBase import FFTBase 5 | 6 | 7 | class T2FFT(FFTBase): 8 | """ 9 | The Fast Fourier Transform on the 2-Torus. 10 | 11 | REMOVE? 12 | 13 | The torus is parameterized by two cyclic variables (x, y). 14 | The standard domain is (x, y) in [0, 1) x [0, 1), in which case the Fourier basis functions are: 15 | exp( i 2 pi xi^T (x; y)) 16 | where xi is the spectral variable, xi in Z^2. 17 | 18 | The Fourier transform is 19 | \hat{f}[p, q] = 1/2pi int_0^2pi f(x, y) exp(-i 2 pi xi^T (x; y)) dx dy 20 | 21 | 22 | 23 | but this class allows one to use arbitrarily scaled and shifted domains D = [l_x, u_x) x [l_y, u_y) 24 | Let the width of the domain be given by 25 | alpha_x = u_x - l_x 26 | alpha_y = u_y - l_y 27 | The basis functions on [l_x, u_x) x [l_y, u_y) are 28 | exp( i 2 pi xi^T ((x - l_x) / alpha_x; (y - l_y) / alpha_y)) 29 | where xi is the spectral variable, xi in Z^2. 30 | The normalized Haar measure is dx dy / (alpha_x * alpha_y) (in terms of Lebesque measure dx dy) 31 | 32 | So the Fourier transform on this particular parameterization of the torus is: 33 | \hat{f}_pq = 1/alpha int_lx^ux int_ly^uy f(x) e^{-2 pi i (p, q)^T ((x - lx) / alpha_x; (y - ly)/alpha_y)} dx dy 34 | 35 | This is what the current class computes, given discrete samples in the domain D. 36 | The samples are assumed to come from the following sampling grid: 37 | (x_i, y_j), i = 0, ... N - 1; j = 0, ..., N - 1 38 | x_i = lx + alpha_x * (i / N_x) 39 | y_i = ly + alpha_y * (i / N_y) 40 | this is the ouput of 41 | x = np.linspace(lx, ux, N_x, endpoint=False) 42 | x = np.linspace(ly, uy, N_y, endpoint=False) 43 | X, Y = np.meshgrid(x, y) 44 | 45 | """ 46 | def __init__(self, lower_bound=(0., 0.), upper_bound=(1., 1.)): 47 | self.lower_bound = np.array(lower_bound) 48 | self.upper_bound = np.array(upper_bound) 49 | 50 | 51 | @staticmethod 52 | def analyze(f, axes=(0, 1)): 53 | """ 54 | Compute the Fourier Transform of the discretely sampled function f : T^2 -> C. 55 | 56 | Let f : T^2 -> C be a band-limited function on the torus. 57 | The samples f(theta_k, phi_l) correspond to points on a regular grid on the circle, 58 | as returned by spaces.T1.linspace: 59 | theta_k = phi_k = 2 pi k / N 60 | for k = 0, ..., N - 1 and l = 0, ..., N - 1 61 | 62 | This function computes 63 | \hat{f}_n = (1/N) \sum_{k=0}^{N-1} f(theta_k) e^{-i n theta_k} 64 | which, if f has band-limit less than N, is equal to: 65 | \hat{f}_n = \int_0^{2pi} f(theta) e^{-i n theta} dtheta / 2pi, 66 | = 67 | where dtheta / 2pi is the normalized Haar measure on T^1, and < , > denotes the inner product on Hilbert space, 68 | with respect to which this transform is unitary. 69 | 70 | The range of frequencies n is -floor(N/2) <= n <= ceil(N/2) - 1 71 | 72 | :param f: 73 | :param axis: 74 | :return: 75 | """ 76 | # The numpy FFT returns coefficients in a different order than we want them, 77 | # and using a different normalization. 78 | f_hat = fft2(f, axes=axes) 79 | f_hat = fftshift(f_hat, axes=axes) 80 | size = np.prod([f.shape[ax] for ax in axes]) 81 | return f_hat / size 82 | 83 | @staticmethod 84 | def synthesize(f_hat, axes=(0, 1)): 85 | """ 86 | :param f_hat: 87 | :param axis: 88 | :return: 89 | """ 90 | 91 | size = np.prod([f_hat.shape[ax] for ax in axes]) 92 | f_hat = ifftshift(f_hat * size, axes=axes) 93 | f = ifft2(f_hat, axes=axes) 94 | return f 95 | -------------------------------------------------------------------------------- /lie_learn/spectral/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/lie_learn/spectral/__init__.py -------------------------------------------------------------------------------- /lie_learn/spectral/fourier_interpolation.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from pynfft import nfft 4 | from pynfft.solver import Solver 5 | 6 | from .T2FFT import T2FFT 7 | 8 | 9 | class FourierInterpolator(object): 10 | 11 | def __init__(self, cartesian_grid_shape, nonequispaced_grid): 12 | """ 13 | The FourierInterpolator can interpolate data on an equispaced Cartesian grid to a non-equispaced grid. 14 | The inpterpolation works by first computing the Fourier coefficients of the input grid, and then evaluating 15 | the Fourier series defined by those coefficients at the non-equispaced grid. 16 | This operation is exactly invertible, as long as the Fourier coefficients are recoverable 17 | from the non-equispaced output samples. 18 | 19 | :param cartesian_grid_shape: the shape (nx, ny) of the input grid. 20 | Samples are assumed to be in [-.5, .5) x [-.5, .5) 21 | :param nonequispaced_grid: the output grid points. Shape (M, 2) 22 | """ 23 | self.cartesian_grid_shape = cartesian_grid_shape 24 | self.nonequispaced_grid_shape = nonequispaced_grid.shape[:-1] 25 | self.nonequispaced_grid = nonequispaced_grid.reshape(-1, 2) 26 | self.nfft = nfft.NFFT(N=cartesian_grid_shape, M=np.prod(nonequispaced_grid.shape[:-1]), 27 | n=None, m=12, flags=None) 28 | self.nfft.x = self.nonequispaced_grid 29 | self.nfft.precompute() 30 | self.solver = Solver(self.nfft) 31 | 32 | 33 | @staticmethod 34 | def init_cartesian_to_polar(nr, nt, nx, ny): 35 | 36 | # On the computation of the polar FFT 37 | # Markus Fenn, Stefan Kunis, Daniel Potts 38 | r = np.linspace(0, 1. / np.sqrt(2), nr) # radius = sqrt((0 - 0.5)^2 + (0 - 0.5)^2) = sqrt(0.5) = 1/sqrt(2) 39 | t = np.linspace(0, 2 * np.pi, nt, endpoint=False) 40 | R, T = np.meshgrid(r, t, indexing='ij') 41 | X = R * np.cos(T) 42 | Y = R * np.sin(T) 43 | C = np.c_[X[..., None], Y[..., None]] 44 | return FourierInterpolator(cartesian_grid_shape=(nx, ny), nonequispaced_grid=C) 45 | 46 | def forward(self, f): 47 | """ 48 | :param f: 49 | :return: 50 | """ 51 | 52 | # Fourier transform x: 53 | # Perform a regular FFT: 54 | f_hat = T2FFT.analyze(f) 55 | 56 | print(f_hat) 57 | 58 | # Since this equispaced FFT assumes spatial samples in theta_k in [0, 1) 59 | # [assuming basis functions exp(i 2 pi n theta), not exp(i n theta)], 60 | # we shift by 0.5, i.e. multiply by exp(-i pi n) = (-1)^n 61 | f_hat *= ((-1) ** np.arange(-np.floor(f.shape[0] / 2), np.ceil(f.shape[0] / 2)))[:, None] 62 | f_hat *= ((-1) ** np.arange(-np.floor(f.shape[1] / 2), np.ceil(f.shape[1] / 2)))[None, :] 63 | 64 | print(f_hat) 65 | 66 | # Use NFFT to evaluate the function defined by these Fourier coefficients at the non-equispaced output grid 67 | self.nfft.f_hat = f_hat 68 | f_resampled = self.nfft.trafo().reshape(self.nonequispaced_grid_shape).copy() 69 | 70 | print(f_resampled) 71 | return f_resampled 72 | 73 | def backward(self, f): 74 | """ 75 | 76 | :param f: 77 | :return: 78 | """ 79 | 80 | self.solver.y = f 81 | self.solver.before_loop() 82 | for i in range(40): 83 | self.solver.loop_one_step() 84 | 85 | f_hat = self.solver.f_hat_iter 86 | 87 | # Since this equispaced FFT assumes spatial samples in theta_k in [0, 1) 88 | # [assuming basis functions exp(i 2 pi n theta), not exp(i n theta)], 89 | # we shift by 0.5, i.e. multiply by exp(-i pi n) = (-1)^n 90 | f_hat /= ((-1) ** np.arange(-np.floor(f_hat.shape[0] / 2), np.ceil(f_hat.shape[0] / 2)))[:, None] 91 | f_hat /= ((-1) ** np.arange(-np.floor(f_hat.shape[1] / 2), np.ceil(f_hat.shape[1] / 2)))[None, :] 92 | f = T2FFT.synthesize(f_hat) 93 | 94 | return f 95 | 96 | 97 | 98 | def test2(): 99 | 100 | nr = 100 101 | nt = 100 102 | nx = 20 103 | ny = 20 104 | 105 | F = FourierInterpolator.init_cartesian_to_polar(nr=nr, nt=nt, nx=nx, ny=ny) 106 | 107 | X, Y = np.meshgrid(np.linspace(-0.5, 0.5, nx, endpoint=False), np.linspace(-0.5, 0.5, ny, endpoint=False), 108 | indexing='ij') 109 | f = np.exp(2*np.pi*1j*(X+0.5)) # + np.exp(2*np.pi * 1j * (3*(Y+0.5))) 110 | C = np.c_[X[..., None], Y[..., None]].reshape(-1, 2) 111 | 112 | #F = FourierInterpolator(grid_in_shape=X.shape, grid_out=C) 113 | print('aa') 114 | fp = F.forward(f) 115 | 116 | fr = F.backward(fp) 117 | return F, f, fp, fr 118 | 119 | 120 | def test(sx=0, sy=0): 121 | 122 | nx = 33 123 | ny = 37 124 | nt = 16 125 | nr = 16 126 | f = np.zeros((nx, ny), dtype='complex') 127 | 128 | F = nfft.NFFT(N=(nx, ny), M=nx * ny) 129 | X, Y = np.meshgrid(np.linspace(-0.5, 0.5, nx, endpoint=False), np.linspace(-0.5, 0.5, ny, endpoint=False), 130 | indexing='ij') 131 | f = np.exp(2*np.pi*1j*(X+0.5)) 132 | F.x = np.c_[X[..., None], Y[..., None]].reshape(-1, 2) 133 | F.precompute() 134 | 135 | f_hat = T2FFT.analyze(f) 136 | tf_hat = f_hat.copy() 137 | tf_hat *= np.exp((2. * np.pi * 1j * sx * np.arange(-np.floor(f.shape[0] / 2.), np.ceil(f.shape[0] / 2.))[:, None]) / f.shape[0]) 138 | tf_hat *= np.exp((2. * np.pi * 1j * sy * np.arange(-np.floor(f.shape[1] / 2.), np.ceil(f.shape[1] / 2.))[None, :]) / f.shape[1]) 139 | 140 | 141 | F.f_hat = f_hat.conj() 142 | f_reconst1 = F.trafo().copy().conj() 143 | 144 | F.f_hat = tf_hat.conj() 145 | f_reconst2 = F.trafo().copy().conj() 146 | 147 | return F, f, f_hat, tf_hat, f_reconst1, f_reconst2 -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = [ 3 | 'setuptools', 4 | 'setuptools-scm', 5 | 'cython', 6 | 'numpy ; python_version>="3.0"', 7 | 'numpy<1.17 ; python_version<"3.0"', 8 | ] 9 | build-backend = "setuptools.build_meta" 10 | 11 | [project] 12 | name = "lie_learn" 13 | version = "0.0.2" 14 | description = "A python package that knows how to do various tricky computations related to Lie groups and manifolds (mainly the sphere S2 and rotation group SO3)." 15 | readme = "README.md" 16 | license = {file = "LICENSE"} 17 | requires-python = ">2.7,!=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, !=3.5.*" 18 | 19 | classifiers = [ 20 | "Programming Language :: Python :: 2.7", 21 | "Programming Language :: Python :: 3", 22 | "Programming Language :: Python :: 3.9", 23 | "Programming Language :: Python :: 3.10", 24 | "Programming Language :: Python :: 3.11", 25 | "Programming Language :: Python :: 3.12", 26 | "License :: OSI Approved :: MIT License", 27 | "Operating System :: OS Independent", 28 | ] 29 | dependencies = [ 30 | 'requests', 31 | 'numpy ; python_version>="3.0"', 32 | 'scipy ; python_version>="3.0"', 33 | 'numpy<1.17 ; python_version<"3.0"', 34 | 'scipy<1.3 ; python_version<"3.0"', 35 | # 'pynfft': # This installation is complicated. Do it yourself. 36 | ] 37 | 38 | [project.urls] 39 | 'Source Code' = "https://github.com/AMLab-Amsterdam/lie_learn" 40 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=missing-docstring 2 | import numpy as np 3 | from Cython.Build import cythonize 4 | from setuptools import setup, find_packages 5 | 6 | setup( 7 | ext_modules=cythonize('lie_learn/**/*.pyx', language_level=2), 8 | include_dirs=[np.get_include()], 9 | ) 10 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/tests/__init__.py -------------------------------------------------------------------------------- /tests/spaces/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/tests/spaces/__init__.py -------------------------------------------------------------------------------- /tests/spaces/test_S3_quadrature.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import lie_learn.spaces.S3 as S3 4 | from lie_learn.representations.SO3.wigner_d import wigner_D_function 5 | 6 | 7 | def test_S3_quadint_equals_numint(): 8 | """Test if SO(3) quadrature integration gives the same result as scipy numerical integration""" 9 | b = 10 10 | for l in range(2): 11 | for m in range(-l, l + 1): 12 | for n in range(-l, l + 1): 13 | check_S3_quadint_equals_numint(l, m, n, b) 14 | 15 | 16 | def check_S3_quadint_equals_numint(l=1, m=1, n=1, b=10): 17 | # Create grids on the sphere 18 | x = S3.meshgrid(b=b, grid_type='SOFT') 19 | x = np.c_[x[0][..., None], x[1][..., None], x[2][..., None]] 20 | 21 | # Compute quadrature weights 22 | w = S3.quadrature_weights(b=b, grid_type='SOFT') 23 | 24 | # Define a polynomial function, to be evaluated at one point or at an array of points 25 | def f1(alpha, beta, gamma): 26 | df = wigner_D_function(l=l, m=m, n=n, alpha=alpha, beta=beta, gamma=gamma) 27 | return df * df.conj() 28 | 29 | def f1a(xs): 30 | d = np.zeros(x.shape[:-1]) 31 | for i in range(d.shape[0]): 32 | for j in range(d.shape[1]): 33 | for k in range(d.shape[2]): 34 | d[i, j, k] = f1(xs[i, j, k, 0], xs[i, j, k, 1], xs[i, j, k, 2]) 35 | return d 36 | 37 | # Obtain the "true" value of the integral of the function over the sphere, using scipy's numerical integration 38 | # routines 39 | i1 = S3.integrate(f1, normalize=True) 40 | 41 | # Compute the integral using the quadrature formulae 42 | i1_w = S3.integrate_quad(f1a(x), grid_type='SOFT', normalize=True, w=w) 43 | 44 | # Check error 45 | print(b, l, m, n, 'results:', i1_w, i1, 'diff:', np.abs(i1_w - i1)) 46 | assert np.isclose(np.abs(i1_w - i1), 0.0) 47 | 48 | -------------------------------------------------------------------------------- /tests/spaces/test_spherical_quadrature.py: -------------------------------------------------------------------------------- 1 | 2 | import lie_learn.spaces.S2 as S2 3 | import numpy as np 4 | 5 | 6 | def test_spherical_quadrature(): 7 | """ 8 | Testing spherical quadrature rule versus numerical integration. 9 | """ 10 | 11 | b = 8 # 10 12 | 13 | # Create grids on the sphere 14 | x_gl = S2.meshgrid(b=b, grid_type='Gauss-Legendre') 15 | x_cc = S2.meshgrid(b=b, grid_type='Clenshaw-Curtis') 16 | x_soft = S2.meshgrid(b=b, grid_type='SOFT') 17 | x_gl = np.c_[x_gl[0][..., None], x_gl[1][..., None]] 18 | x_cc = np.c_[x_cc[0][..., None], x_cc[1][..., None]] 19 | x_soft = np.c_[x_soft[0][..., None], x_soft[1][..., None]] 20 | 21 | # Compute quadrature weights 22 | w_gl = S2.quadrature_weights(b=b, grid_type='Gauss-Legendre') 23 | w_cc = S2.quadrature_weights(b=b, grid_type='Clenshaw-Curtis') 24 | w_soft = S2.quadrature_weights(b=b, grid_type='SOFT') 25 | 26 | # Define a polynomial function, to be evaluated at one point or at an array of points 27 | def f1a(xs): 28 | xc = S2.change_coordinates(coords=xs, p_from='S', p_to='C') 29 | return xc[..., 0] ** 2 * xc[..., 1] - 1.4 * xc[..., 2] * xc[..., 1] ** 3 + xc[..., 1] - xc[..., 2] ** 2 + 2. 30 | def f1(theta, phi): 31 | xs = np.array([theta, phi]) 32 | return f1a(xs) 33 | 34 | # Obtain the "true" value of the integral of the function over the sphere, using scipy's numerical integration 35 | # routines 36 | i1 = S2.integrate(f1, normalize=False) 37 | 38 | # Compute the integral using the quadrature formulae 39 | # i1_gl_w = (w_gl * f1a(x_gl)).sum() 40 | i1_gl_w = S2.integrate_quad(f1a(x_gl), grid_type='Gauss-Legendre', normalize=False, w=w_gl) 41 | print(i1_gl_w, i1, 'diff:', np.abs(i1_gl_w - i1)) 42 | assert np.isclose(np.abs(i1_gl_w - i1), 0.0) 43 | 44 | # i1_cc_w = (w_cc * f1a(x_cc)).sum() 45 | i1_cc_w = S2.integrate_quad(f1a(x_cc), grid_type='Clenshaw-Curtis', normalize=False, w=w_cc) 46 | print(i1_cc_w, i1, 'diff:', np.abs(i1_cc_w - i1)) 47 | assert np.isclose(np.abs(i1_cc_w - i1), 0.0) 48 | 49 | i1_soft_w = (w_soft * f1a(x_soft)).sum() 50 | print(i1_soft_w, i1, 'diff:', np.abs(i1_soft_w - i1)) 51 | print(i1_soft_w) 52 | print(i1) 53 | # assert np.isclose(np.abs(i1_cc_w - i1), 0.0) # TODO 54 | -------------------------------------------------------------------------------- /tests/spectral/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AMLab-Amsterdam/lie_learn/edf012f5f60af320175d2e6269db78b984b5bfc3/tests/spectral/__init__.py -------------------------------------------------------------------------------- /tests/spectral/test_S2FFT.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import lie_learn.spaces.S2 as S2 4 | from lie_learn.representations.SO3.spherical_harmonics import sh 5 | from lie_learn.spectral.S2FFT import setup_legendre_transform, setup_legendre_transform_indices, sphere_fft, S2_FT_Naive 6 | 7 | 8 | def test_S2_FT_Naive(): 9 | 10 | L_max = 6 11 | 12 | for grid_type in ('Gauss-Legendre', 'Clenshaw-Curtis'): 13 | 14 | theta, phi = S2.meshgrid(b=L_max + 1, grid_type=grid_type) 15 | 16 | for field in ('real', 'complex'): 17 | for normalization in ('quantum', 'seismology'): # TODO Others should work but are not normalized 18 | for condon_shortley in ('cs', 'nocs'): 19 | 20 | fft = S2_FT_Naive(L_max, grid_type=grid_type, 21 | field=field, normalization=normalization, condon_shortley=condon_shortley) 22 | 23 | for l in range(L_max): 24 | for m in range(-l, l + 1): 25 | 26 | y_true = sh( 27 | l, m, theta, phi, 28 | field=field, normalization=normalization, condon_shortley=condon_shortley == 'cs') 29 | 30 | y_hat = fft.analyze(y_true) 31 | 32 | # The flat index for (l, m) is l^2 + l + m 33 | # Before the harmonics of degree l, there are this many harmonics: 34 | # sum_{i=0}^{l-1} 2i+1 = l^2 35 | # There are 2l+1 harmonics of degree l, with order m=0 at the center, 36 | # so the m-th harmonic of degree is at l + m within the block of degree l. 37 | y_hat_true = np.zeros_like(y_hat) 38 | y_hat_true[l**2 + l + m] = 1 39 | 40 | y = fft.synthesize(y_hat_true) 41 | 42 | diff = np.sum(np.abs(y_hat - y_hat_true)) 43 | print(grid_type, field, normalization, condon_shortley, l, m, diff) 44 | assert np.isclose(diff, 0.) 45 | 46 | diff = np.sum(np.abs(y - y_true)) 47 | print(grid_type, field, normalization, condon_shortley, l, m, diff) 48 | assert np.isclose(diff, 0.) 49 | 50 | 51 | def test_S2FFT(): 52 | 53 | L_max = 10 54 | beta, alpha = S2.meshgrid(b=L_max + 1, grid_type='Driscoll-Healy') 55 | lt = setup_legendre_transform(b=L_max + 1) 56 | lti = setup_legendre_transform_indices(b=L_max + 1) 57 | 58 | for l in range(L_max): 59 | for m in range(-l, l + 1): 60 | 61 | Y = sh(l, m, beta, alpha, 62 | field='complex', normalization='seismology', condon_shortley=True) 63 | 64 | y_hat = sphere_fft(Y, lt, lti) 65 | 66 | # The flat index for (l, m) is l^2 + l + m 67 | # Before the harmonics of degree l, there are this many harmonics: sum_{i=0}^{l-1} 2i+1 = l^2 68 | # There are 2l+1 harmonics of degree l, with order m=0 at the center, 69 | # so the m-th harmonic of degree is at l + m within the block of degree l. 70 | y_hat_true = np.zeros_like(y_hat) 71 | y_hat_true[l**2 + l + m] = 1 72 | 73 | diff = np.sum(np.abs(y_hat - y_hat_true)) 74 | nz = 1. - np.isclose(y_hat, 0.) 75 | diff_nz = np.sum(np.abs(nz - y_hat_true)) 76 | print(l, m, diff, diff_nz) 77 | print(np.round(y_hat, 4)) 78 | print(y_hat_true) 79 | # assert np.isclose(diff, 0.) # TODO make this work 80 | print(nz) 81 | assert np.isclose(diff_nz, 0.) 82 | 83 | -------------------------------------------------------------------------------- /tests/spectral/test_S2FFT_NFFT.py: -------------------------------------------------------------------------------- 1 | import lie_learn.spaces.S2 as S2 2 | from lie_learn.spectral.S2FFT_NFFT import S2FFT_NFFT 3 | from lie_learn.representations.SO3.spherical_harmonics import * 4 | 5 | 6 | def test_S2FFT_NFFT(): 7 | """ 8 | Testing S2FFT NFFT 9 | """ 10 | b = 8 11 | convention = 'Gauss-Legendre' 12 | #convention = 'Clenshaw-Curtis' 13 | x = S2.meshgrid(b=b, grid_type=convention) 14 | print(x[0].shape, x[1].shape) 15 | x = np.c_[x[0][..., None], x[1][..., None]]#.reshape(-1, 2) 16 | print(x.shape) 17 | x = x.reshape(-1, 2) 18 | w = S2.quadrature_weights(b=b, grid_type=convention).flatten() 19 | F = S2FFT_NFFT(L_max=b, x=x, w=w) 20 | 21 | for l in range(0, b): 22 | for m in range(-l, l + 1): 23 | #l = b; m = b 24 | f = sh(l, m, x[..., 0], x[..., 1], field='real', normalization='quantum', condon_shortley=True) 25 | #f2 = np.random.randn(*f.shape) 26 | print(f) 27 | 28 | f_hat = F.analyze(f) 29 | print(np.round(f_hat, 3)) 30 | f_reconst = F.synthesize(f_hat) 31 | 32 | #print np.round(f, 3) 33 | print(np.round(f_reconst, 3)) 34 | #print np.round(f/f_reconst, 3) 35 | print(np.abs(f-f_reconst).sum()) 36 | assert np.isclose(np.abs(f-f_reconst).sum(), 0.) 37 | 38 | print(np.round(f_hat, 3)) 39 | assert np.isclose(f_hat[l ** 2 + l + m], 1.) 40 | #assert False -------------------------------------------------------------------------------- /tests/spectral/test_S2_conv.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | 4 | import lie_learn.spaces.S2 as S2 5 | import lie_learn.spaces.S3 as S3 6 | import lie_learn.groups.SO3 as SO3 7 | from lie_learn.representations.SO3.spherical_harmonics import sh 8 | from lie_learn.spectral.S2_conv import naive_S2_conv, spectral_S2_conv, naive_S2_conv_v2 9 | 10 | 11 | def compare_naive_and_spectral_conv(): 12 | 13 | f1 = lambda t, p: sh(l=2, m=1, theta=t, phi=p, field='real', normalization='quantum', condon_shortley=True) 14 | f2 = lambda t, p: sh(l=2, m=1, theta=t, phi=p, field='real', normalization='quantum', condon_shortley=True) 15 | 16 | theta, phi = S2.meshgrid(b=4, grid_type='Gauss-Legendre') 17 | f1_grid = f1(theta, phi) 18 | f2_grid = f2(theta, phi) 19 | 20 | alpha, beta, gamma = S3.meshgrid(b=4, grid_type='SOFT') # TODO check convention 21 | 22 | f12_grid_spectral = spectral_S2_conv(f1_grid, f2_grid, s2_fft=None, so3_fft=None) 23 | 24 | f12_grid = np.zeros_like(alpha) 25 | for i in range(alpha.shape[0]): 26 | for j in range(alpha.shape[1]): 27 | for k in range(alpha.shape[2]): 28 | f12_grid[i, j, k] = naive_S2_conv(f1, f2, alpha[i, j, k], beta[i, j, k], gamma[i, j, k]) 29 | print(i, j, k, f12_grid[i, j, k]) 30 | 31 | return f1_grid, f2_grid, f12_grid, f12_grid_spectral 32 | 33 | 34 | def naive_conv(l1=1, m1=1, l2=1, m2=1, g_parameterization='EA313'): 35 | f1 = lambda t, p: sh(l=l1, m=m1, theta=t, phi=p, field='real', normalization='quantum', condon_shortley=True) 36 | f2 = lambda t, p: sh(l=l2, m=m2, theta=t, phi=p, field='real', normalization='quantum', condon_shortley=True) 37 | 38 | theta, phi = S2.meshgrid(b=3, grid_type='Gauss-Legendre') 39 | f1_grid = f1(theta, phi) 40 | f2_grid = f2(theta, phi) 41 | 42 | alpha, beta, gamma = S3.meshgrid(b=3, grid_type='SOFT') # TODO check convention 43 | 44 | f12_grid = np.zeros_like(alpha) 45 | for i in range(alpha.shape[0]): 46 | for j in range(alpha.shape[1]): 47 | for k in range(alpha.shape[2]): 48 | f12_grid[i, j, k] = naive_S2_conv_v2(f1, f2, alpha[i, j, k], beta[i, j, k], gamma[i, j, k], g_parameterization) 49 | print(i, j, k, f12_grid[i, j, k]) 50 | 51 | return f1_grid, f2_grid, f12_grid 52 | -------------------------------------------------------------------------------- /tests/spectral/test_SO3_FFT_Naive.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | from lie_learn.spectral.SO3FFT_Naive import SO3_FFT_NaiveReal, SO3_FFT_SemiNaive_Complex, SO3_FT_Naive 4 | from lie_learn.representations.SO3.pinchon_hoggan.pinchon_hoggan_dense import Jd, rot_mat 5 | from lie_learn.representations.SO3.irrep_bases import change_of_basis_matrix 6 | 7 | 8 | # TODO: test if the Fourier transform of a right SO(2)-invariant function is zero except for a column at n=0, and 9 | # test if it is equal to the spherical harmonics transform of the corresponding function on the sphere 10 | 11 | 12 | def test_SO3_FT_Naive(): 13 | """ 14 | Check that the naive complex SO(3) FFT: 15 | - Produces the right Wigner-D function when given a 1-hot input to the synthesis transform 16 | - Produces a 1-hot vector when given a single Wigner-D function to the analysis transform 17 | """ 18 | L_max = 3 19 | 20 | f_hat = [np.zeros((2 * ll + 1, 2 * ll + 1)) for ll in range(L_max + 1)] 21 | 22 | # TODO: the SO3_FFT_SemiNaive_Complex no longer uses the D convention parameters because of new caching feature 23 | 24 | field = 'complex' 25 | order = 'centered' 26 | for normalization in ('quantum', 'seismology'): # Note: the geodesy and nfft wigners are normalized differently 27 | for condon_shortley in ('cs', 'nocs'): 28 | 29 | fft = SO3_FT_Naive(L_max=L_max, 30 | field=field, normalization=normalization, 31 | order=order, condon_shortley=condon_shortley) 32 | 33 | for l in range(L_max + 1): 34 | for m in range(-l, l + 1): 35 | for n in range(-l, l + 1): 36 | f_hat[l][l + m, l + n] = 1. / (2 * l + 1) 37 | f_hat_flat = np.hstack([fhl.flatten() for fhl in f_hat]) 38 | D = fft.synthesize_by_matmul(f_hat_flat) 39 | 40 | D2 = make_D_sample_grid(b=L_max + 1, l=l, m=m, n=n, 41 | field=field, normalization=normalization, 42 | order=order, condon_shortley=condon_shortley) 43 | 44 | diff = np.sum(np.abs(D - D2.flatten())) 45 | print(l, m, n, 'Synthesize error:', diff) 46 | assert np.isclose(diff, 0.0) 47 | 48 | f_hat_2 = fft.analyze_by_matmul(D2) 49 | 50 | # f_hat_flat = np.hstack([ff.flatten() for ff in f_hat]) 51 | f_hat_2_flat = np.hstack([ff.flatten() for ff in f_hat_2]) 52 | 53 | # f_hat_2_flat *= (2 * l + 1) # / (4 * np.pi) # apply magic constant TODO fix this 54 | print(f_hat_2_flat) 55 | print(f_hat_flat) 56 | print(np.max(np.abs(f_hat_flat)), np.max(np.abs(f_hat_2_flat))) 57 | 58 | diff = np.sum(np.abs(f_hat_flat - f_hat_2_flat)) 59 | print(l, m, n, 'Analyze error:', diff) 60 | assert np.isclose(diff, 0.0) 61 | 62 | f_hat[l][l + m, l + n] = 0. 63 | 64 | def test_SO3_FFT_SemiNaiveComplex(): 65 | """ 66 | Check that the naive complex SO(3) FFT: 67 | - Produces the right Wigner-D function when given a 1-hot input to the synthesis transform 68 | - Produces a 1-hot vector when given a single Wigner-D function to the analysis transform 69 | """ 70 | L_max = 3 71 | 72 | f_hat = [np.zeros((2 * ll + 1, 2 * ll + 1)) for ll in range(L_max + 1)] 73 | 74 | # TODO: the SO3_FFT_SemiNaive_Complex no longer uses the D convention parameters because of new caching feature 75 | 76 | field = 'complex' 77 | order = 'centered' 78 | for normalization in ('quantum', 'seismology'): # Note: the geodesy and nfft wigners are normalized differently 79 | for condon_shortley in ('cs', 'nocs'): 80 | 81 | fft = SO3_FFT_SemiNaive_Complex(L_max=L_max, L2_normalized=False, 82 | field=field, normalization=normalization, 83 | order=order, condon_shortley=condon_shortley) 84 | 85 | #fft = SO3_FFT_Naive(L_max=L_max, 86 | # field=field, normalization=normalization, 87 | # order=order, condon_shortley=condon_shortley) 88 | 89 | for l in range(L_max + 1): 90 | for m in range(-l, l + 1): 91 | for n in range(-l, l + 1): 92 | f_hat[l][l + m, l + n] = 1. 93 | D = fft.synthesize(f_hat) 94 | 95 | D2 = make_D_sample_grid(b=L_max + 1, l=l, m=m, n=n, 96 | field=field, normalization=normalization, 97 | order=order, condon_shortley=condon_shortley) 98 | 99 | diff = np.sum(np.abs(D - D2)) 100 | print(l, m, n, diff) 101 | assert np.isclose(diff, 0.0) 102 | 103 | f_hat_2 = fft.analyze(D2) 104 | 105 | f_hat_flat = np.hstack([ff.flatten() for ff in f_hat]) 106 | f_hat_2_flat = np.hstack([ff.flatten() for ff in f_hat_2]) 107 | 108 | f_hat_2_flat *= (2 * l + 1) / (4 * np.pi) # apply magic constant TODO fix this 109 | 110 | diff = np.sum(np.abs(f_hat_flat - f_hat_2_flat)) 111 | print(l, m, n, diff) 112 | assert np.isclose(diff, 0.0) 113 | 114 | f_hat[l][l + m, l + n] = 0. 115 | 116 | 117 | # TODO: test linearity of FFT 118 | 119 | #TODO 120 | def check_SO3_FFT_NaiveComplex_invertible(): 121 | L_max = 3 122 | 123 | f_hat = [np.zeros((2 * ll + 1, 2 * ll + 1)) for ll in range(L_max + 1)] 124 | fft = SO3_FFT_SemiNaive_Complex(L_max=L_max, L2_normalized=False) 125 | for l in range(L_max + 1): 126 | for m in range(-l, l + 1): 127 | for n in range(-l, l + 1): 128 | f_hat[l][l + m, l + n] = 1. 129 | f = fft.synthesize(f_hat) 130 | f_hat_2 = fft.analyze(f) 131 | 132 | diff = np.sum([np.abs(f_hat[ll] - f_hat_2[ll]) for ll in range(L_max + 1)]) 133 | 134 | f_hat[l][l + m, l + n] = 0. 135 | print(l, m, n, diff) # , D2 / D 136 | assert np.isclose(diff, 0.0) 137 | 138 | 139 | def test_SO3_FFT_NaiveReal(): 140 | """ 141 | Testing if the real Naive SO(3) FFT synthesis works correctly for 1-hot input vectors 142 | """ 143 | L_max = 3 144 | 145 | f_hat = [np.zeros((2 * ll + 1, 2 * ll + 1)) for ll in range(L_max + 1)] 146 | fft = SO3_FFT_NaiveReal(L_max=L_max, L2_normalized=False) 147 | for l in range(L_max + 1): 148 | for m in range(-l, l + 1): 149 | for n in range(-l, l + 1): 150 | f_hat[l][l + m, l + n] = 1. 151 | D = fft.synthesize(f_hat) 152 | f_hat[l][l + m, l + n] = 0. 153 | D2 = make_D_sample_grid(b=L_max + 1, l=l, m=m, n=n, 154 | field='real', normalization='quantum', order='centered', condon_shortley='cs') 155 | 156 | print(l, m, n, np.sum(np.abs(D - D2))) 157 | assert np.isclose(np.sum(np.abs(D - D2)), 0.0) 158 | 159 | 160 | def make_D_sample_grid(b=4, l=0, m=0, n=0, 161 | field='complex', normalization='seismology', order='centered', condon_shortley='cs'): 162 | 163 | from lie_learn.representations.SO3.wigner_d import wigner_D_function 164 | D = lambda a, b, c: wigner_D_function(l, m, n, alpha, beta, gamma, 165 | field=field, normalization=normalization, 166 | order=order, condon_shortley=condon_shortley) 167 | 168 | f = np.zeros((2 * b, 2 * b, 2 * b), dtype='complex') 169 | 170 | for j1 in range(f.shape[0]): 171 | alpha = 2 * np.pi * j1 / (2. * b) 172 | for k in range(f.shape[1]): 173 | beta = np.pi * (2 * k + 1) / (4. * b) 174 | for j2 in range(f.shape[2]): 175 | gamma = 2 * np.pi * j2 / (2. * b) 176 | f[j1, k, j2] = D(alpha, beta, gamma) 177 | return f 178 | -------------------------------------------------------------------------------- /tests/spectral/test_conv_S2_SO3.py: -------------------------------------------------------------------------------- 1 | 2 | import lie_learn.spaces.S2 as S2 3 | from lie_learn.spaces.S3 import change_coordinates 4 | 5 | 6 | 7 | 8 | 9 | --------------------------------------------------------------------------------