microlux 0.1.0__tar.gz

microlux-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Haibin Ren
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

microlux-0.1.0/PKG-INFO

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: microlux
+Version: 0.1.0
+Summary: A JAX-based package for microlensing binary lens
+Author-email: Haibin Ren <rhb23@mails.tsinghua.edu.cn>
+License: MIT License
+        
+        Copyright (c) 2024 Haibin Ren
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Project-URL: repository, https://github.com/CoastEgo/microlux
+Keywords: jax,microlensing,binary lens
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.4.28
+Requires-Dist: numpy
+Requires-Dist: scipy
+Dynamic: license-file
+
+## microlux: Microlensing using Jax
+
+[![Test Status](https://github.com/coastego/microlux/actions/workflows/run_test.yml/badge.svg)](https://github.com/CoastEgo/microlux/actions/workflows/run_test.yml)
+[![Documentation Status](https://github.com/coastego/microlux/actions/workflows/build_docs.yml/badge.svg)](https://coastego.github.io/microlux/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+---
+
+`microlux` is a <a href='https://github.com/jax-ml/jax'>Jax</a>-based package that can calculate the binary lensing light curve and its derivatives both efficiently and accurately.  We use the modified adaptive contour integratoin in <a href='https://github.com/valboz/VBBinaryLensing'>`VBBinaryLensing`</a> to maximize the performance. 
+With the access to the gradient, we can use more advanced algorithms for microlensing modeling, such as Hamiltonian Monte Carlo (HMC) in <a href='https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.hmc.NUTS'>`numpyro`</a>.
+
+
+## Installation
+
+``` bash
+pip install microlux
+```
+or you can install this package from source for development. 
+``` bash 
+git clone https://github.com/CoastEgo/microlux.git
+cd microlux
+pip install -e .
+```
+## Documentation
+The documentation is available at <a href='https://coastego.github.io/microlux/'>here</a>. See this for more details.
+
+
+  
+## Citation
+
+`microlux` is open-source software licensed under the MIT license. If you use this package for your research, please cite our paper:
+
+- A differentiable binary microlensing model using adaptive contour integration method: <a href='https://arxiv.org/abs/2501.07268'>in arXiv</a> and <a href= 'https://iopscience.iop.org/article/10.3847/1538-3881/adb1b2'>in AJ </a>.
+
+``` bibtex
+@article{ren2025microlux,
+       author = {{Ren}, Haibin and {Zhu}, Wei},
+        title = "{A Differentiable Binary Microlensing Model Using Adaptive Contour Integration Method}",
+      journal = {The Astronomical Journal},
+         year = 2025,
+       volume = {169},
+       number = {3},
+          eid = {170},
+        pages = {170},
+          doi = {10.3847/1538-3881/adb1b2},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2025AJ....169..170R},
+}
+```

microlux-0.1.0/README.md

@@ -0,0 +1,47 @@
+## microlux: Microlensing using Jax
+
+[![Test Status](https://github.com/coastego/microlux/actions/workflows/run_test.yml/badge.svg)](https://github.com/CoastEgo/microlux/actions/workflows/run_test.yml)
+[![Documentation Status](https://github.com/coastego/microlux/actions/workflows/build_docs.yml/badge.svg)](https://coastego.github.io/microlux/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+---
+
+`microlux` is a <a href='https://github.com/jax-ml/jax'>Jax</a>-based package that can calculate the binary lensing light curve and its derivatives both efficiently and accurately.  We use the modified adaptive contour integratoin in <a href='https://github.com/valboz/VBBinaryLensing'>`VBBinaryLensing`</a> to maximize the performance. 
+With the access to the gradient, we can use more advanced algorithms for microlensing modeling, such as Hamiltonian Monte Carlo (HMC) in <a href='https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.hmc.NUTS'>`numpyro`</a>.
+
+
+## Installation
+
+``` bash
+pip install microlux
+```
+or you can install this package from source for development. 
+``` bash 
+git clone https://github.com/CoastEgo/microlux.git
+cd microlux
+pip install -e .
+```
+## Documentation
+The documentation is available at <a href='https://coastego.github.io/microlux/'>here</a>. See this for more details.
+
+
+  
+## Citation
+
+`microlux` is open-source software licensed under the MIT license. If you use this package for your research, please cite our paper:
+
+- A differentiable binary microlensing model using adaptive contour integration method: <a href='https://arxiv.org/abs/2501.07268'>in arXiv</a> and <a href= 'https://iopscience.iop.org/article/10.3847/1538-3881/adb1b2'>in AJ </a>.
+
+``` bibtex
+@article{ren2025microlux,
+       author = {{Ren}, Haibin and {Zhu}, Wei},
+        title = "{A Differentiable Binary Microlensing Model Using Adaptive Contour Integration Method}",
+      journal = {The Astronomical Journal},
+         year = 2025,
+       volume = {169},
+       number = {3},
+          eid = {170},
+        pages = {170},
+          doi = {10.3847/1538-3881/adb1b2},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2025AJ....169..170R},
+}
+```

microlux-0.1.0/pyproject.toml

@@ -0,0 +1,45 @@
+[project]
+name = "microlux"
+version = "0.1.0"
+description = "A JAX-based package for microlensing binary lens"
+authors = [
+    {name = "Haibin Ren", email = "rhb23@mails.tsinghua.edu.cn"},
+]
+requires-python = ">=3.9"
+license = {file = "LICENSE"}
+readme = "README.md"
+urls = {repository = "https://github.com/CoastEgo/microlux"}
+keywords = ["jax", "microlensing", "binary lens"]
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+]
+dependencies = [
+    "jax >=0.4.28",
+    "numpy",
+    "scipy",
+]
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["microlux"]
+
+[tool.ruff]
+extend-include = ["*.ipynb"]
+src = []
+
+[tool.ruff.lint]
+fixable = ["I001", "F401"]
+ignore = ["E501","E402", "E721", "E731", "E741", "F722","E702"]
+select = ["E", "F", "I001"]
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+extra-standard-library = ["typing_extensions"]
+lines-after-imports = 2
+order-by-type = false

microlux-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

microlux-0.1.0/src/microlux/__init__.py

@@ -0,0 +1,26 @@
+# # -*- coding: utf-8 -*-
+all = [
+    "point_light_curve",
+    "extended_light_curve",
+    "contour_integral",
+    "binary_mag",
+    "Iterative_State",
+    "Error_State",
+    "to_lowmass",
+    "to_centroid",
+]
+
+from .basic_function import (
+    to_centroid as to_centroid,
+    to_lowmass as to_lowmass,
+)
+from .countour import contour_integral as contour_integral
+from .model import (
+    binary_mag as binary_mag,
+    extended_light_curve as extended_light_curve,
+    point_light_curve as point_light_curve,
+)
+from .utils import (
+    Error_State as Error_State,
+    Iterative_State as Iterative_State,
+)

microlux-0.1.0/src/microlux/basic_function.py

@@ -0,0 +1,236 @@
+import jax
+import jax.numpy as jnp
+
+
+jax.config.update("jax_enable_x64", True)
+
+
+def to_centroid(s, q, x):
+    """
+    Transforms the coordinate system to the centroid.
+
+    Parameters:
+    s (float): The projected separation between the two objects.
+    q (float): The planet to host mass ratio.
+    x (complex): The original coordinate.
+
+    Returns:
+    complex: The transformed coordinate in the centroid system.
+    """
+    delta_x = s / (1 + q)
+    return -(jnp.conj(x) - delta_x)
+
+
+def to_lowmass(s, q, x):
+    """
+    Transforms the coordinate system to the system where the lower mass object is at the origin.
+
+    Parameters:
+    s (float): The separation between the two components.
+    q (float): The mass ratio of the two components.
+    x (complex): The original centroid coordinate.
+
+    Returns:
+    complex: The transformed coordinate in the low mass component coordinate system.
+    """
+    delta_x = s / (1 + q)
+    return -jnp.conj(x) + delta_x
+
+
+def Quadrupole_test(rho, s, q, zeta, z, cond, tol=1e-2):
+    """
+    The quadrupole test, ghost image test, and planetary caustic test proposed by Bozza 2010 to check the validity of the point source approximation.
+    The coefficients are fine-tuned in our implementation.
+
+    """
+    m1 = 1 / (1 + q)
+    m2 = q / (1 + q)
+    cQ = 2
+    cG = 3
+    cP = 4  # tunable parameters vbbl 2018 + version=3.6.2 choose cQ=3,cG=miu_G (vbbl typo) ,cP=4
+
+    # basic derivatives
+    fz0 = lambda z: -m1 / (z - s) - m2 / z
+    fz1 = lambda z: m1 / (z - s) ** 2 + m2 / z**2
+    fz2 = lambda z: -2 * m1 / (z - s) ** 3 - 2 * m2 / z**3
+    fz3 = lambda z: 6 * m1 / (z - s) ** 4 + 6 * m2 / z**4
+    J = lambda z: 1 - fz1(z) * jnp.conj(fz1(z))
+
+    ####Quadrupole test
+    miu_Q = jnp.abs(
+        -2
+        * jnp.real(
+            3 * jnp.conj(fz1(z)) ** 3 * fz2(z) ** 2
+            - (3 - 3 * J(z) + J(z) ** 2 / 2) * jnp.abs(fz2(z)) ** 2
+            + J(z) * jnp.conj(fz1(z)) ** 2 * fz3(z)
+        )
+        / (J(z) ** 5)
+    )
+
+    # cusp test
+    miu_C = jnp.abs(jnp.imag(3 * jnp.conj(fz1(z)) ** 3 * fz2(z) ** 2) / (J(z) ** 5))
+    mag = jnp.sum(jnp.where(cond, jnp.abs(1 / J(z)), 0), axis=1)
+    cond1 = (
+        jnp.sum(jnp.where(cond, (miu_Q + miu_C), 0), axis=1)
+        * cQ
+        * (rho**2 + 1e-4 * tol)
+        < tol
+    )
+
+    ####ghost image test
+    zwave = jnp.conj(zeta) - fz0(z)
+    J_wave = 1 - fz1(z) * fz1(zwave)
+    J3 = J_wave * fz2(jnp.conj(z)) * fz1(z)
+    miu_G = jnp.abs((J3 - jnp.conj(J3) * fz1(zwave)) / (J(z) * J_wave**2))
+    miu_G = jnp.where(cond, 0, miu_G)
+    cond2 = ((rho + 1e-3) * miu_G * cG < 1).all(axis=1)  # all() is same with VBBL code
+
+    #####planet test # in our frame primary is at s, the planet is at 0, so the position of the planetary caustic is 1/s
+    cond3 = (
+        (q > 1e-2)
+        | (
+            jnp.abs(zeta - 1 / s) ** 2
+            > cP * (rho**2 + 9 * q / s**2)  # rho**2*s**2<q comment out in vbbl 3.6.2
+        )
+    )[:, 0]
+
+    return cond1 & cond2 & cond3, mag
+
+
+def get_poly_coff(zeta_l, s, m2):
+    """
+    get the polynomial cofficients of the polynomial equation of the lens equation. The low mass object is at the origin and the primary is at s.
+    The input zeta_l should have the shape of (n,1) for broadcasting.
+    """
+    zeta_conj = jnp.conj(zeta_l)
+    c0 = s**2 * zeta_l * m2**2
+    c1 = -s * m2 * (2 * zeta_l + s * (-1 + s * zeta_l - 2 * zeta_l * zeta_conj + m2))
+    c2 = (
+        zeta_l
+        - s**3 * zeta_l * zeta_conj
+        + s * (-1 + m2 - 2 * zeta_conj * zeta_l * (1 + m2))
+        + s**2 * (zeta_conj - 2 * zeta_conj * m2 + zeta_l * (1 + zeta_conj**2 + m2))
+    )
+    c3 = (
+        s**3 * zeta_conj
+        + 2 * zeta_l * zeta_conj
+        + s**2 * (-1 + 2 * zeta_conj * zeta_l - zeta_conj**2 + m2)
+        - s * (zeta_l + 2 * zeta_l * zeta_conj**2 - 2 * zeta_conj * m2)
+    )
+    c4 = zeta_conj * (-1 + 2 * s * zeta_conj + zeta_conj * zeta_l) - s * (
+        -1 + 2 * s * zeta_conj + zeta_conj * zeta_l + m2
+    )
+    c5 = (s - zeta_conj) * zeta_conj
+    coff = jnp.concatenate((c5, c4, c3, c2, c1, c0), axis=1)
+    return coff
+
+
+def get_zeta_l(rho, trajectory_centroid_l, theta):  # 获得等高线采样的zeta
+    zeta_l = trajectory_centroid_l + rho * jnp.exp(1j * theta)
+    return zeta_l
+
+
+def verify(zeta_l, z_l, s, m1, m2):  # verify whether the root is right
+    return jnp.abs(z_l - m1 / (jnp.conj(z_l) - s) - m2 / jnp.conj(z_l) - zeta_l)
+
+
+def get_parity(z, s, m1, m2):  # get the parity of roots
+    de_conjzeta_z1 = m1 / (jnp.conj(z) - s) ** 2 + m2 / jnp.conj(z) ** 2
+    return jnp.sign((1 - jnp.abs(de_conjzeta_z1) ** 2))
+
+
+def get_parity_error(z, s, m1, m2):
+    de_conjzeta_z1 = m1 / (jnp.conj(z) - s) ** 2 + m2 / jnp.conj(z) ** 2
+    return jnp.abs((1 - jnp.abs(de_conjzeta_z1) ** 2))
+
+
+def dot_product(a, b):
+    return jnp.real(a) * jnp.real(b) + jnp.imag(a) * jnp.imag(b)
+
+
+def basic_partial(z, theta, rho, q, s, caustic_crossing):
+    """
+
+    basic partial derivatives of the lens equation with respect to zeta, z, and theta used in the error estimation.
+
+    """
+    z_c = jnp.conj(z)
+    parZetaConZ = 1 / (1 + q) * (1 / (z_c - s) ** 2 + q / z_c**2)
+    par2ConZetaZ = -2 / (1 + q) * (1 / (z - s) ** 3 + q / (z) ** 3)
+    de_zeta = 1j * rho * jnp.exp(1j * theta)
+    detJ = 1 - jnp.abs(parZetaConZ) ** 2
+    de_z = (de_zeta - parZetaConZ * jnp.conj(de_zeta)) / detJ
+    deXProde2X = (rho**2 + jnp.imag(de_z**2 * de_zeta * par2ConZetaZ)) / detJ
+
+    def get_de_deXPro_de2X(carry):
+        # now only calculate the derivative of x'^x'' with respect to \theta if caustic_crossing is True which is used in e4 calculation
+        # still need to test weather this is robust enough for the case that source is very close to the caustic but not crossing it
+
+        de2_zeta = -rho * jnp.exp(1j * theta)
+        de2_zetaConj = -rho * jnp.exp(-1j * theta)
+        par3ConZetaZ = 6 / (1 + q) * (1 / (z - s) ** 4 + q / (z) ** 4)
+        de2_z = (
+            de2_zeta
+            - jnp.conj(par2ConZetaZ) * jnp.conj(de_z) ** 2
+            - parZetaConZ * (de2_zetaConj - par2ConZetaZ * de_z**2)
+        ) / detJ
+        # deXProde2X_test = 1/(2*1j)*(de2_z*jnp.conj(de_z)-de_z*jnp.conj(de2_z))
+        # jax.debug.print('deXProde2X_test error is {}',jnp.nansum(jnp.abs(deXProde2X_test-deXProde2X)))
+        de_deXPro_de2X = (
+            1
+            / detJ**2
+            * jnp.imag(
+                detJ
+                * (
+                    de2_zeta * par2ConZetaZ * de_z**2
+                    + de_zeta * par3ConZetaZ * de_z**3
+                    + de_zeta * par2ConZetaZ * 2 * de_z * de2_z
+                )
+                + (
+                    jnp.conj(par2ConZetaZ) * jnp.conj(de_z) * jnp.conj(parZetaConZ)
+                    + parZetaConZ * par2ConZetaZ * de_z
+                )
+                * de_zeta
+                * par2ConZetaZ
+                * de_z**2
+            )
+        )
+        return de_deXPro_de2X
+
+    de_deXPro_de2X = jax.lax.cond(
+        caustic_crossing, get_de_deXPro_de2X, lambda x: jnp.zeros_like(deXProde2X), None
+    )
+    # deXProde2X = jax.lax.stop_gradient(deXProde2X)
+    return deXProde2X, de_z, de_deXPro_de2X
+
+
+@jax.custom_jvp
+def refine_gradient(zeta_l, q, s, z):
+    return z
+
+
+@refine_gradient.defjvp
+def refine_gradient_jvp(primals, tangents):
+    """
+    use the custom jvp to refine the gradient of roots respect to zeta_l, based on the equation on V.Bozza 2010 eq 20 and also see our paper for the details
+    This will simplify the computational graph and accelerate the gradient calculation
+    """
+    zeta, q, s, z = primals
+    tangent_zeta, tangent_q, tangent_s, tangent_z = tangents
+
+    z_c = jnp.conj(z)
+    parZetaConZ = 1 / (1 + q) * (1 / (z_c - s) ** 2 + q / z_c**2)
+    detJ = 1 - jnp.abs(parZetaConZ) ** 2
+
+    parZetaq = 1 / (1 + q) ** 2 * (1 / (z_c - s) - 1 / z_c)
+    add_item_q = tangent_q * (parZetaq - jnp.conj(parZetaq) * parZetaConZ)
+
+    parZetas = -1 / (1 + q) / (z_c - s) ** 2
+    add_item_s = tangent_s * (parZetas - jnp.conj(parZetas) * parZetaConZ)
+
+    tangent_z2 = (
+        tangent_zeta - parZetaConZ * jnp.conj(tangent_zeta) - add_item_q - add_item_s
+    ) / detJ
+    # tangent_z2 = jnp.where(jnp.isnan(tangent_z2),0.,tangent_z2)
+    # jax.debug.print('{}',(tangent_z2-tangent_z).sum())
+    return z, tangent_z2