@@ -35,12 +35,10 @@ def check_dep(f, *args, **kwargs):
3535
3636@timer
3737def test_gsparams ():
38- if is_jax_galsim ():
39- pass
40- else :
41- check_dep (galsim .GSParams , allowed_flux_variation = 0.90 )
42- check_dep (galsim .GSParams , range_division_for_extrema = 50 )
43- check_dep (galsim .GSParams , small_fraction_of_flux = 1.e-6 )
38+ if is_jax_galsim (): return
39+ check_dep (galsim .GSParams , allowed_flux_variation = 0.90 )
40+ check_dep (galsim .GSParams , range_division_for_extrema = 50 )
41+ check_dep (galsim .GSParams , small_fraction_of_flux = 1.e-6 )
4442
4543
4644@timer
@@ -498,6 +496,8 @@ def test_hsm_depr():
498496
499497@timer
500498def test_photon_array_depr ():
499+ if is_jax_galsim (): return
500+
501501 nphotons = 1000
502502
503503 # First create from scratch
@@ -544,203 +544,96 @@ def test_photon_array_depr():
544544
545545 # Using the getter is allowed, but deprecated.
546546 photon_array = galsim .PhotonArray (nphotons )
547- if is_jax_galsim ():
548- # jax-galsim always sets these additional properties
549- # dxdz = check_dep(getattr, photon_array, 'dxdz')
550- # however jax-galsim sets them to NaN so they are not allocated
551- assert not photon_array .hasAllocatedAngles ()
552- assert len (photon_array .dxdz ) == nphotons
553- # JAX-Galsim does not allow by reference setting - changed this
554- # to make tests below run
555- photon_array .dxdz = 0.17
556- # non-nan means allocated for jax-galsim
557- assert photon_array .hasAllocatedAngles ()
558- else :
559- dxdz = check_dep (getattr , photon_array , 'dxdz' )
560- assert photon_array .hasAllocatedAngles ()
561- assert photon_array .hasAllocatedAngles ()
562- assert len (photon_array .dxdz ) == nphotons
563- assert len (photon_array .dydz ) == nphotons
564- dxdz [:] = 0.17
547+ dxdz = check_dep (getattr , photon_array , 'dxdz' )
548+ assert photon_array .hasAllocatedAngles ()
549+ assert photon_array .hasAllocatedAngles ()
550+ assert len (photon_array .dxdz ) == nphotons
551+ assert len (photon_array .dydz ) == nphotons
552+ dxdz [:] = 0.17
565553 np .testing .assert_array_equal (photon_array .dxdz , 0.17 )
566554 np .testing .assert_array_equal (photon_array .dydz , 0. )
567555
568- if is_jax_galsim ():
569- assert hasattr (photon_array , "dydz" )
570- # JAX-Galsim does not allow by reference setting - changed this
571- # to make tests below run
572- photon_array .dydz = 0.59
573- else :
574- dydz = photon_array .dydz # Allowed now.
575- dydz [:] = 0.59
556+ dydz = photon_array .dydz # Allowed now.
557+ dydz [:] = 0.59
576558 np .testing .assert_array_equal (photon_array .dydz , 0.59 )
577559
578- if is_jax_galsim ():
579- # jax-galsim always sets these additional properties
580- # wave = check_dep(getattr, photon_array, 'wavelength')
581- # however jax-galsim sets them to NaN so they are not allocated
582- assert not photon_array .hasAllocatedWavelengths ()
583- else :
584- wave = check_dep (getattr , photon_array , 'wavelength' )
585- assert photon_array .hasAllocatedWavelengths ()
560+ wave = check_dep (getattr , photon_array , 'wavelength' )
561+ assert photon_array .hasAllocatedWavelengths ()
586562 assert len (photon_array .wavelength ) == nphotons
587- if is_jax_galsim ():
588- # JAX-Galsim does not allow by reference setting - changed this
589- # to make tests below run
590- photon_array .wavelength = 500.0
591- # jax-galsim is allocated now
592- assert photon_array .hasAllocatedWavelengths ()
593- else :
594- wave [:] = 500.
563+ wave [:] = 500.
595564 np .testing .assert_array_equal (photon_array .wavelength , 500 )
596565
597- if is_jax_galsim ():
598- # jax-galsim always sets these additional properties
599- # u = check_dep(getattr, photon_array, "pupil_u")
600- # however jax-galsim sets them to NaN so they are not allocated
601- assert not photon_array .hasAllocatedPupil ()
602- else :
603- u = check_dep (getattr , photon_array , 'pupil_u' )
604- assert photon_array .hasAllocatedPupil ()
605- assert len (photon_array .pupil_u ) == nphotons
566+ u = check_dep (getattr , photon_array , 'pupil_u' )
567+ assert photon_array .hasAllocatedPupil ()
568+ assert len (photon_array .pupil_u ) == nphotons
606569 assert len (photon_array .pupil_v ) == nphotons
607- if is_jax_galsim ():
608- # JAX-Galsim does not allow by reference setting - changed this
609- # to make tests below run
610- photon_array .pupil_u = 6.0
611- # jax-galsim is allocated now
612- assert photon_array .hasAllocatedPupil ()
613- else :
614- u [:] = 6.0
570+ u [:] = 6.0
615571 np .testing .assert_array_equal (photon_array .pupil_u , 6.0 )
616572 np .testing .assert_array_equal (photon_array .pupil_v , 0.0 )
617- if is_jax_galsim ():
618- assert hasattr (photon_array , "pupil_v" )
619- # JAX-Galsim does not allow by reference setting - changed this
620- # to make tests below run
621- photon_array .pupil_v = 10.0
622- else :
623- v = photon_array .pupil_v
624- v [:] = 10.0
573+ v = photon_array .pupil_v
574+ v [:] = 10.0
625575 np .testing .assert_array_equal (photon_array .pupil_v , 10.0 )
626576
627- if is_jax_galsim ():
628- # jax-galsim always sets these additional properties
629- # t = check_dep(getattr, photon_array, "time")
630- # however jax-galsim sets them to NaN so they are not allocated
631- assert not photon_array .hasAllocatedTimes ()
632- # jax-galsim needs to set 0
633- photon_array .time = 0.0
634- # jax-galsim is allocated now
635- assert photon_array .hasAllocatedTimes ()
636- else :
637- t = check_dep (getattr , photon_array , 'time' )
577+ t = check_dep (getattr , photon_array , 'time' )
638578 assert photon_array .hasAllocatedTimes ()
639579 assert len (photon_array .time ) == nphotons
640580 np .testing .assert_array_equal (photon_array .time , 0.0 )
641- if is_jax_galsim ():
642- # JAX-Galsim does not allow by reference setting - changed this
643- # to make tests below run
644- photon_array .time = 10
645- else :
646- t [:] = 10
581+ t [:] = 10
647582 np .testing .assert_array_equal (photon_array .time , 10.0 )
648583
649584 # For coverage, also need to test the two pair ones in other order.
650585 photon_array = galsim .PhotonArray (nphotons )
651- if is_jax_galsim ():
652- # jax-galsim always sets these additional properties
653- # dydz = check_dep(getattr, photon_array, "dydz")
654- # however jax-galsim sets them to NaN so they are not allocated
655- assert not photon_array .hasAllocatedAngles ()
656- else :
657- dydz = check_dep (getattr , photon_array , 'dydz' )
658- assert photon_array .hasAllocatedAngles ()
659- assert photon_array .hasAllocatedAngles ()
586+ dydz = check_dep (getattr , photon_array , 'dydz' )
587+ assert photon_array .hasAllocatedAngles ()
588+ assert photon_array .hasAllocatedAngles ()
660589 assert len (photon_array .dxdz ) == nphotons
661590 assert len (photon_array .dydz ) == nphotons
662- if is_jax_galsim ():
663- # JAX-Galsim does not allow by reference setting - changed this
664- # to make tests below run
665- photon_array .dydz = 0.59
666- # non-nan means allocated for jax-galsim
667- assert photon_array .hasAllocatedAngles ()
668- else :
669- dydz [:] = 0.59
591+ dydz [:] = 0.59
670592 np .testing .assert_array_equal (photon_array .dxdz , 0. )
671593 np .testing .assert_array_equal (photon_array .dydz , 0.59 )
672594
673- if is_jax_galsim ():
674- assert hasattr (photon_array , "dxdz" )
675- # JAX-Galsim does not allow by reference setting - changed this
676- # to make tests below run
677- photon_array .dxdz = 0.17
678- else :
679- dxdz = photon_array .dxdz # Allowed now.
680- dxdz [:] = 0.17
595+ dxdz = photon_array .dxdz # Allowed now.
596+ dxdz [:] = 0.17
681597 np .testing .assert_array_equal (photon_array .dxdz , 0.17 )
682598
683- if is_jax_galsim ():
684- # jax-galsim always sets these additional properties
685- # v = check_dep(getattr, photon_array, "pupil_v")
686- # however jax-galsim sets them to NaN so they are not allocated
687- assert not photon_array .hasAllocatedPupil ()
688- else :
689- v = check_dep (getattr , photon_array , 'pupil_v' )
690- assert photon_array .hasAllocatedPupil ()
691- assert len (photon_array .pupil_u ) == nphotons
599+ v = check_dep (getattr , photon_array , 'pupil_v' )
600+ assert photon_array .hasAllocatedPupil ()
601+ assert len (photon_array .pupil_u ) == nphotons
692602 assert len (photon_array .pupil_v ) == nphotons
693- if is_jax_galsim ():
694- # JAX-Galsim does not allow by reference setting - changed this
695- # to make tests below run
696- photon_array .pupil_v = 10.0
697- # jax-galsim is allocated now
698- assert photon_array .hasAllocatedPupil ()
699- else :
700- v [:] = 10.0
603+ v [:] = 10.0
701604 np .testing .assert_array_equal (photon_array .pupil_u , 0.0 )
702605 np .testing .assert_array_equal (photon_array .pupil_v , 10.0 )
703- if is_jax_galsim ():
704- assert hasattr (photon_array , "pupil_u" )
705- # JAX-Galsim does not allow by reference setting - changed this
706- # to make tests below run
707- photon_array .pupil_u = 6.0
708- else :
709- u = photon_array .pupil_u
710- u [:] = 6.0
606+ u = photon_array .pupil_u
607+ u [:] = 6.0
711608 np .testing .assert_array_equal (photon_array .pupil_u , 6.0 )
712609
713610 # Check assignAt
714611 pa1 = galsim .PhotonArray (50 )
715612 pa1 .x = photon_array .x [:50 ]
716- if is_jax_galsim ():
717- pa1 .y = photon_array .y [:50 ]
718- pa1 .flux .at [0 :50 ].set (photon_array .flux [:50 ])
719- else :
720- for i in range (50 ):
721- pa1 .y [i ] = photon_array .y [i ]
722- pa1 .flux [0 :50 ] = photon_array .flux [:50 ]
613+ for i in range (50 ):
614+ pa1 .y [i ] = photon_array .y [i ]
615+ pa1 .flux [0 :50 ] = photon_array .flux [:50 ]
723616 pa1 .dxdz = photon_array .dxdz [:50 ]
724617 pa1 .dydz = photon_array .dydz [:50 ]
725618 pa1 .pupil_u = photon_array .pupil_u [:50 ]
726619 pa1 .pupil_v = photon_array .pupil_v [:50 ]
727620 pa2 = galsim .PhotonArray (100 )
728621 check_dep (pa2 .assignAt , 0 , pa1 )
729622 check_dep (pa2 .assignAt , 50 , pa1 )
730- np .testing .assert_array_almost_equal (pa2 .x [:50 ], pa1 .x )
731- np .testing .assert_array_almost_equal (pa2 .y [:50 ], pa1 .y )
732- np .testing .assert_array_almost_equal (pa2 .flux [:50 ], pa1 .flux )
733- np .testing .assert_array_almost_equal (pa2 .dxdz [:50 ], pa1 .dxdz )
734- np .testing .assert_array_almost_equal (pa2 .dydz [:50 ], pa1 .dydz )
735- np .testing .assert_array_almost_equal (pa2 .pupil_u [:50 ], pa1 .pupil_u )
736- np .testing .assert_array_almost_equal (pa2 .pupil_v [:50 ], pa1 .pupil_v )
737- np .testing .assert_array_almost_equal (pa2 .x [50 :], pa1 .x )
738- np .testing .assert_array_almost_equal (pa2 .y [50 :], pa1 .y )
739- np .testing .assert_array_almost_equal (pa2 .flux [50 :], pa1 .flux )
740- np .testing .assert_array_almost_equal (pa2 .dxdz [50 :], pa1 .dxdz )
741- np .testing .assert_array_almost_equal (pa2 .dydz [50 :], pa1 .dydz )
742- np .testing .assert_array_almost_equal (pa2 .pupil_u [50 :], pa1 .pupil_u )
743- np .testing .assert_array_almost_equal (pa2 .pupil_v [50 :], pa1 .pupil_v )
623+ np .testing .assert_almost_equal (pa2 .x [:50 ], pa1 .x )
624+ np .testing .assert_almost_equal (pa2 .y [:50 ], pa1 .y )
625+ np .testing .assert_almost_equal (pa2 .flux [:50 ], pa1 .flux )
626+ np .testing .assert_almost_equal (pa2 .dxdz [:50 ], pa1 .dxdz )
627+ np .testing .assert_almost_equal (pa2 .dydz [:50 ], pa1 .dydz )
628+ np .testing .assert_almost_equal (pa2 .pupil_u [:50 ], pa1 .pupil_u )
629+ np .testing .assert_almost_equal (pa2 .pupil_v [:50 ], pa1 .pupil_v )
630+ np .testing .assert_almost_equal (pa2 .x [50 :], pa1 .x )
631+ np .testing .assert_almost_equal (pa2 .y [50 :], pa1 .y )
632+ np .testing .assert_almost_equal (pa2 .flux [50 :], pa1 .flux )
633+ np .testing .assert_almost_equal (pa2 .dxdz [50 :], pa1 .dxdz )
634+ np .testing .assert_almost_equal (pa2 .dydz [50 :], pa1 .dydz )
635+ np .testing .assert_almost_equal (pa2 .pupil_u [50 :], pa1 .pupil_u )
636+ np .testing .assert_almost_equal (pa2 .pupil_v [50 :], pa1 .pupil_v )
744637
745638 # Error if it doesn't fit.
746639 with assert_raises (ValueError ):
0 commit comments