Jax to Main!#188

Merged

Jammy2211 merged 240 commits intomainfrom

feature/jax_wrapper

Oct 6, 2025

Commits on Apr 4, 2025

Commits on Apr 6, 2025

Commits on Apr 7, 2025

Commits on Apr 8, 2025

Commits on Apr 9, 2025

Commits on Apr 21, 2025

moved over most zoom-y stuff

Jammy2211
authored and
Jammy2211
committed
removed more zoom stuf

Jammy2211
authored and
Jammy2211
committed
more stuff moved over

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
updated to use new zoom API

Jammy2211
authored and
Jammy2211
committed

Commits on Apr 29, 2025

zoom refactor done

Jammy2211
authored and
Jammy2211
committed
remove some tests

Jammy2211
authored and
Jammy2211
committed

Commits on May 4, 2025

update extracted_array_2d_from to not use numba

Jammy2211
authored and
Jammy2211
committed
updat eunit test shapes

Jammy2211
authored and
Jammy2211
committed
unit test fix

Jammy2211
authored and
Jammy2211
committed
fix visualization output

Jammy2211
authored and
Jammy2211
committed
unit test fixes

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 11, 2025

black and minor fixes

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 12, 2025

updated factory test on two mappers to use more stable solver

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 13, 2025

fix test__identical_inversion_source_and_image_loops

Jammy2211
authored and
Jammy2211
committed
check for nans when raising InversionException

Jammy2211
authored and
Jammy2211
committed
test__inversion_imaging__via_linear_obj_func_and_mapper__force_edge_pixels_to_zero

Jammy2211
authored and
Jammy2211
committed
fix test__inversion_imaging__linear_obj_func_with_w_tilde

Jammy2211
authored and
Jammy2211
committed
fix test__inversion_imaging__linear_obj_func_and_non_func_give_same_terms

Jammy2211
authored and
Jammy2211
committed
fix test__inversion_imaging__linear_obj_func_with_w_tilde

Jammy2211
authored and
Jammy2211
committed
test__identical_inversion_values_for_two_methods

Jammy2211
authored and
Jammy2211
committed
midway through some refactoring

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 15, 2025

remove force_edge_image_pixels_to_zeros

Jammy2211
authored and
Jammy2211
committed
convert JAX Array

Jammy2211
authored and
Jammy2211
committed
ndarray conversions for JAX arrays

Jammy2211
authored and
Jammy2211
committed
fix all inversion unit tests

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
Merge pull request #175 from Jammy2211/feature/jax_linear_light
Jammy2211
authored

Commits on Jun 17, 2025

removed numpy based Trianghle and fixed first test

Jammy2211
authored and
Jammy2211
committed
fix one_triangle.triangles

Jammy2211
authored and
Jammy2211
committed
test__above

Jammy2211
authored and
Jammy2211
committed
multiple fixes in test_coordinate_implementation.py

Jammy2211
authored and
Jammy2211
committed
remove containment test

Jammy2211
authored and
Jammy2211
committed
test_array_representation deleted

Jammy2211
authored and
Jammy2211
committed
fix test_extended_source

Jammy2211
authored and
Jammy2211
committed
JAx CoordinateArrayTriangles has explicit JAX use now

Jammy2211
authored and
Jammy2211
committed
simplfiied coordinate_array and removed support for numpy

Jammy2211
authored and
Jammy2211
committed
refactor arrayu to only use JAX

Jammy2211
authored and
Jammy2211
committed
simplify unit tests

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
update to much simpler tests for transformer

Jammy2211
authored and
Jammy2211
committed
remove pylops legact

Jammy2211
authored and
Jammy2211
committed
fix unitt est with shaping

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 18, 2025

preload_real_transforms converted to numpy

Jammy2211
authored and
Jammy2211
committed
preload_imag_transforms

Jammy2211
authored and
Jammy2211
committed
rename functions

Jammy2211
authored and
Jammy2211
committed
transformed_mapping_matrix_from

Jammy2211
authored and
Jammy2211
committed
transformed_mapping_matrix_via_preload_from

Jammy2211
authored and
Jammy2211
committed
simplify Transformer

Jammy2211
authored and
Jammy2211
committed
TransformerDFT docstring

Jammy2211
authored and
Jammy2211
committed
transformer now uses jax arrays in DFT

Jammy2211
authored and
Jammy2211
committed
remove ordered_1d

Jammy2211
authored and
Jammy2211
committed
JAX intereferometer grad works

Jammy2211
authored and
Jammy2211
committed
dft_preload_transform added to Interferometer inputs

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
Merge pull request #176 from Jammy2211/feature/jax_only_triangles
Jammy2211
authored
refactor use of dirty image in inversion

Jammy2211
authored and
Jammy2211
committed
JAX on transforfed mapping matrtix

Jammy2211
authored and
Jammy2211
committed
update plot

Jammy2211
authored and
Jammy2211
committed
Merge pull request #177 from Jammy2211/feature/jax_interferometer
Jammy2211
authored
resized_array_2d_from converted to numpy

Jammy2211
authored and
Jammy2211
committed
remove replace_noise_map_2d_values_where_image_2d_values_are_negative

Jammy2211
authored and
Jammy2211
committed
index_2d_for_index_slim_from no longer uses numba

Jammy2211
authored and
Jammy2211
committed
index_slim_for_index_2d_from does not use numba

Jammy2211
authored and
Jammy2211
committed
remove array_2d_slim_complex_from

Jammy2211
authored and
Jammy2211
committed
remove array_2d_native_complex_via_indexes_from

Jammy2211
authored and
Jammy2211
committed
remove numba from two more functions ing rid_2d_util

Jammy2211
authored and
Jammy2211
committed
remove grid_2d_slim_upscaled_from

Jammy2211
authored and
Jammy2211
committed
remove native_sub_index_for_slim_sub_index_2d_from

Jammy2211
authored and
Jammy2211
committed
remove redudnant tests

Jammy2211
authored and
Jammy2211
committed
slim_index_for_sub_slim_index_via_mask_2d_from no longer uses numba

Jammy2211
authored and
Jammy2211
committed
removed sub_slim_index_for_sub_native_index_from

Jammy2211
authored and
Jammy2211
committed
remove oversample_mask_2d_from

Jammy2211
authored and
Jammy2211
committed
sub_size_radial_bins_from no longer uses numba

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 19, 2025

grid_2d_slim_over_sampled_via_mask_from does not use mask

Jammy2211
authored and
Jammy2211
committed
removed final numba functions not used for inversion

Jammy2211
authored and
Jammy2211
committed
numba stops pixelization if not installed

Jammy2211
authored and
Jammy2211
committed
test

Jammy2211
authored and
Jammy2211
committed
merge

Jammy2211
authored and
Jammy2211
committed
use simpler grid_2d_slim_over_sampled_via_mask_from which works

Jammy2211
authored and
Jammy2211
committed
temporary solution

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 23, 2025

remove profile func

Jammy2211
authored and
Jammy2211
committed
remove run_time_dict

Jammy2211
authored and
Jammy2211
committed
remove run time dict

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 24, 2025

removed all profilng and tested

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
Merge pull request #179 from Jammy2211/feature/jax_remove_profiling
Jammy2211
authored
rectangular uses intterpolation with JAX support now

Jammy2211
authored and
Jammy2211
committed
fix visualization unit tests

Jammy2211
authored and
Jammy2211
committed
test_autoarray/inversion/pixelization/mappers/test_rectangular.py

Jammy2211
authored and
Jammy2211
committed
test_autoarray/inversion/pixelization/mappers/test_factory.py -> rectanguilar mapping matrix now has interpoltion

Jammy2211
authored and
Jammy2211
committed
interpolate works but now need to remove convolver

Jammy2211
authored and
Jammy2211
committed
convolver convler mapping matrix replaced for psf, fixes issue

Jammy2211
authored and
Jammy2211
committed
fix test__inversion_matrices__x2_mappers

Jammy2211
authored and
Jammy2211
committed
fix test_autoarray/inversion/inversion/imaging/test_imaging.py

Jammy2211
authored and
Jammy2211
committed
fix test__curvature_matrix

Jammy2211
authored and
Jammy2211
committed
fix test__inversion_imaging__via_mapper

Jammy2211
authored and
Jammy2211
committed
all unit tests pass

Jammy2211
authored and
Jammy2211
committed
Merge pull request #181 from Jammy2211/feature/jax_remove_convolver
Jammy2211
authored
clean up imports

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 25, 2025

comment out test

Jammy2211
authored and
Jammy2211
committed
fix fit _util whilst reudcing imports

Jammy2211
authored and
Jammy2211
committed
remove many np.array() conversions

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
remove cached properties from BorderRelocator for easier JAx

Jammy2211
authored and
Jammy2211
committed
removed for loop from border function

Jammy2211
authored and
Jammy2211
committed
converted relocated_grid_from to use JAX

Jammy2211
authored and
Jammy2211
committed
border relocator function converred to JAX

Jammy2211
authored and
Jammy2211
committed
grid_2d_slim_via_shape_native_not_mask_from

Jammy2211
authored and
Jammy2211
committed
Rectangular fidxes

Jammy2211
authored and
Jammy2211
committed

Commits on Jun 26, 2025

fix w tild eiwth some ndarray conversions

Jammy2211
authored and
Jammy2211
committed
fix over sampling tests

Jammy2211
authored and
Jammy2211
committed
fix plotting

Jammy2211
authored and
Jammy2211
committed
fix some more tests due to numba jax

Jammy2211
authored and
Jammy2211
committed
coment out test to get past it for now, think its just linear lagebra stability

Jammy2211
authored and
Jammy2211
committed
updated _Reducedd matrices to use zeroing

Jammy2211
authored and
Jammy2211
committed
regularization_matrix_Reduced

Jammy2211
authored and
Jammy2211
committed
fix test

Jammy2211
authored and
Jammy2211
committed
add preloading in order to pass mapper indexes

Jammy2211
authored and
Jammy2211
committed
full JAX success

Jammy2211
authored and
Jammy2211
committed
adaptive_pixel_signals_from JAX-d

Jammy2211
authored and
Jammy2211
committed
convert mapped_to_source_via_mapping_matrix_from to numpy

Jammy2211
authored and
Jammy2211
committed
update data_weight_total_for_pix_from

Jammy2211
authored and
Jammy2211
committed
moved sub_slim_indexes_for_pix_index to inversion_interferometer_util

Jammy2211
authored and
Jammy2211
committed

Commits on Jul 12, 2025

Merge pull request #182 from Jammy2211/feature/jax_speed_up_general
Jammy2211
authored

Commits on Jul 13, 2025

convert soem regularization util functions from numba to numpy

Jammy2211
authored and
Jammy2211
committed
remove minus ones in data_weight_total_for_pix_from

Jammy2211
authored and
Jammy2211
committed
mapper index list returns ndarray

Jammy2211
authored and
Jammy2211
committed
fix autoarray/inversion/inversion/interferometer/w_tilde.py

Jammy2211
authored and
Jammy2211
committed
fixed bug where regularization matrix was returned for curvature_reg_reduced

Jammy2211
authored and
Jammy2211
committed
fix case where border relocator is off

Jammy2211
authored and
Jammy2211
committed
w tilde now default to false

Jammy2211
authored and
Jammy2211
committed
remove old source pixel zeroing functionality

Jammy2211
authored and
Jammy2211
committed
docuemnet preloasds and mapper_index_list -> mapper_indices

Jammy2211
authored and
Jammy2211
committed
fix last unit test

Jammy2211
authored and
Jammy2211
committed
black

Jammy2211
authored and
Jammy2211
committed
Merge pull request #183 from Jammy2211/feature/jax_inversion
Jammy2211
authored

Commits on Jul 14, 2025

removed all reference to include in plotting

Jammy2211
authored and
Jammy2211
committed

Commits on Jul 21, 2025

visual clean up complete

Jammy2211
authored and
Jammy2211
committed

Commits on Jul 22, 2025

Merge pull request #184 from Jammy2211/feature/jax_simplify_visualization
Jammy2211
authored

Commits on Jul 23, 2025

regularization util JAX conversons, seem to work

Jammy2211
authored and
Jammy2211
committed
gaussian kernel converted successfully

Jammy2211
authored and
Jammy2211
committed
fix constant zeroth

Jammy2211
authored and
Jammy2211
committed
move utils to their specific modules

Jammy2211
authored and
Jammy2211
committed
regulsirztion refactor complete and rectangular works

Jammy2211
authored and
Jammy2211
committed

Commits on Jul 24, 2025

updates for rectangular grid sorted

Jammy2211
authored and
Jammy2211
committed
Merge pull request #185 from Jammy2211/feature/jax_rectangular_slam
Jammy2211
authored

Commits on Jul 28, 2025

fix pos neg

Jammy2211
authored and
Jammy2211
committed

Commits on Aug 11, 2025

refactor kernel mapping for speed

Jammy2211
authored and
Jammy2211
committed

Commits on Aug 18, 2025

backup

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 15, 2025

adpative stuff implemented

Jammy2211
authored and
Jammy2211
committed
edges unit test passes

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 16, 2025

rectangular edges

Jammy2211
authored and
Jammy2211
committed
edges_transformed

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 18, 2025

remove InversionException tests

Jammy2211
authored and
Jammy2211
committed
comment out edges test

Jammy2211
authored and
Jammy2211
committed
data_vector_via_w_tilde_data_imaging_from converted to JAX

Jammy2211
authored and
Jammy2211
committed
_data_vector_mapper in JAX

Jammy2211
authored and
Jammy2211
committed
_curvature_matrix_x1_mapper

Jammy2211
authored and
Jammy2211
committed
mapped_reconstructed_data_via_w_tilde_from

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 21, 2025

lots of hacky stuff

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 22, 2025

remove commented code

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 23, 2025

fixed mappings and weights

Jammy2211
authored and
Jammy2211
committed
test__edges_transformed

Jammy2211
authored and
Jammy2211
committed
test__areas_transformed

Jammy2211
authored and
Jammy2211
committed
all of inversion tests pass

Jammy2211
authored and
Jammy2211
committed
all unit tests pass

Jammy2211
authored and
Jammy2211
committed
Merge pull request #186 from Jammy2211/feature/jax_wrapper_rectangular_adaptive
Jammy2211
authored

Commits on Sep 25, 2025

inversion_util

Jammy2211
authored and
Jammy2211
committed
back to stable

Jammy2211
authored and
Jammy2211
committed

Commits on Sep 30, 2025

inned supeer sampling array uses jax.ops.segment_sum for fast compile

Jammy2211
authored and
Jammy2211
committed

Commits on Oct 1, 2025

remove offset subtraction

Jammy2211
authored and
Jammy2211
committed
moved first 3 numba functions

Jammy2211
authored and
Jammy2211
committed
moved more functions

Jammy2211
authored and
Jammy2211
committed
update dataset w tilde to use numba util

Jammy2211
authored and
Jammy2211
committed
update w_tilde to use numba utils

Jammy2211
authored and
Jammy2211
committed
fix _data_vector_x1_mapper

Jammy2211
authored and
Jammy2211
committed
more direct split of numba code

Jammy2211
authored and
Jammy2211
committed
fix unitt est due to jax operated mapping matrix

Jammy2211
authored and
Jammy2211
committed
test__data_linear_func_matrix_dict

Jammy2211
authored and
Jammy2211
committed
update data_linear_func_matrix_from to not use numba

Jammy2211
authored and
Jammy2211
committed
remove use of frames in off diag

Jammy2211
authored and
Jammy2211
committed
docstring

Jammy2211
authored and
Jammy2211
committed
remove convolver

Jammy2211
authored and
Jammy2211
committed
split mappeR_util with numba

Jammy2211
authored and
Jammy2211
committed
split mesh_util into numba

Jammy2211
authored and
Jammy2211
committed
replace more unit tests so they use numba util

Jammy2211
authored and
Jammy2211
committed
fix mesh numab improt

Jammy2211
authored and
Jammy2211
committed

Commits on Oct 6, 2025

JAX complet

Jammy2211
authored and
Jammy2211
committed
Merge pull request #187 from Jammy2211/feature/jax_numba_mix
Jammy2211
authored