Implement JAX automatic differentiation on Turboflow

[RoberAgro]

• [x] Mathematics functions
• [x] Velocity triangle functions
• [x] Cascade functions
• [x] Loss model functions
  • [x] Isentropic (start with this one)
  • [x] Benner loss model
• [x] Deviation model functions
• [x] Overall turbine calculation
• [x] Performance Analysis
• [x] Design Optimization

[SrinivasDiwanji]

Verified the gradients for all the cascade functions and also for evaluate_axial_turbine

Educated guess still required to get meaningful derivative values for comparison

[SrinivasDiwanji]

Pandas dataframe usage has been removed in performance analysis. It has been replaced by dictionaries and other functions decribed in numerics and file_utils docs

[SrinivasDiwanji]

Writing and storing the results gives errors as it uses pandas commands. This has to be updated and changed.

[SrinivasDiwanji]

Jax works with isentropic and critical mach choking criteira for the complete axial turbine problem

[SrinivasDiwanji]

The issue is with critical mass flow rate choking criteria for which jax is not working. Jax differentiates compute_critical_values function in choking criteria but Jax fails when differentiating critical_mass_flow_rate function

[SrinivasDiwanji]

Jax runs for whole turbine calculation with critical_mach_number and kacker-okapuu loss model

[SrinivasDiwanji]

Jax runs for whole turbine calculation with critical_mach_number and moustapha loss model

[SrinivasDiwanji]

Now JAX runs for whole turbine calculation with critical_mach_number and all loss models

[SrinivasDiwanji]

JAX works for two_stage turbine calculation too with crtical_mach_number as choking criterion and all loss models. Now the issue is with critical_mass_flow_rate choking criterion, it has a tracing error with JAX. One possible issue is with how the function is defined. Checking the issue

[SrinivasDiwanji]

JAX works now with critical_mass_flow_rate choking criterion too. The tracer error was due to the fact that critical_state was initialized within the critical_mass_flow_rate function and it was being updated which caused the tracer error while using JAX. It is resolved by using another dict updated_critical_state to avoid mutation of critical_state. The updated code works now

[SrinivasDiwanji]

Performance analysis differentiation with JAX has been done

[SrinivasDiwanji]

Design Optimization works with critical mach number and approx derivatives

[SrinivasDiwanji]

Some probable issues with design_optimization:

1. In evaluate_constraints a jax array is initialized and then being updated which can cause a problem
2. Similar in-place modifications are being done in fitness_gradient_based for geometry, equality constraints array which can cause a problem
3. Strings stator and rotor cannot form a part of jax array and can cause a problem

[SrinivasDiwanji]

len() doesn't work with jax array