From f2cf47a547b2e8015e057334b05236b64e66e0f6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 18 Jun 2025 15:29:12 +0100 Subject: [PATCH 1/2] operate image now supports JAX --- autogalaxy/operate/image.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py index da457eb8c..602d3dcbe 100644 --- a/autogalaxy/operate/image.py +++ b/autogalaxy/operate/image.py @@ -1,6 +1,6 @@ from __future__ import annotations +import jax import jax.numpy as jnp -import numpy as np from typing import TYPE_CHECKING, Dict, List, Optional from autoarray import Array2D @@ -10,7 +10,6 @@ import autoarray as aa -from autogalaxy import exc class OperateImage: @@ -189,12 +188,14 @@ def visibilities_from( image_2d = self.image_2d_from(grid=grid) - if not jnp.any(image_2d.array): - return aa.Visibilities.zeros( + return jax.lax.cond( + jnp.any(image_2d.array), + lambda _: transformer.visibilities_from(image=image_2d), + lambda _: aa.Visibilities.zeros( shape_slim=(transformer.uv_wavelengths.shape[0],) - ) - - return transformer.visibilities_from(image=image_2d) + ), + operand=None + ) class OperateImageList(OperateImage): From 6575e9a1829852e1efcbb4e3ab262bf14e22feee Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 18 Jun 2025 15:41:58 +0100 Subject: [PATCH 2/2] black --- autogalaxy/operate/image.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/autogalaxy/operate/image.py b/autogalaxy/operate/image.py index 602d3dcbe..44a3ab9c6 100644 --- a/autogalaxy/operate/image.py +++ b/autogalaxy/operate/image.py @@ -11,7 +11,6 @@ import autoarray as aa - class OperateImage: """ Packages methods which operate on the 2D image returned from the `image_2d_from` function of a light object @@ -194,7 +193,7 @@ def visibilities_from( lambda _: aa.Visibilities.zeros( shape_slim=(transformer.uv_wavelengths.shape[0],) ), - operand=None + operand=None, )