TensorFlow error when compiling power spectrum function

[EiffL]

While trying to build a tf.function to compute the lensing power spectrum from the simulation, we are running into the following error:

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: Const_1:0

Below is a MWE to reproduce the issue:

import tensorflow as tf
import numpy as np
from flowpm.tfbackground import cosmo,z2a,a2z,afactor, chifactor
from flowpm.tfpower import linear_matter_power
from flowpm.angular_power_tf import measure_power_spectrum_tf
from flowpm.raytracing import  lightcone, Born
import flowpm
from flowpm.tfpower import linear_matter_power

nc=[32,32,320]   # size of the cube, number of cells
plane_size=32                    # number of pixel for x and  y 
Boxsize=[200,200,2000]          # Physical size of the cube
r = np.linspace(0,2000,10, endpoint=True)
a = afactor(r)  
a_s=z2a(1.00)
ds=chifactor(a_s)
field=5.
a0=0.1
af=1.0
n_steps=10

# This is the function that we are trying to compile
@tf.function
def power_spectrum_for_cosmology(
              Omega0_m,
              sigma8):
    cosmology=cosmo
    cosmology['Omega0_m']=tf.convert_to_tensor(Omega0_m,dtype=tf.float32)
    cosmology['sigma8']=tf.convert_to_tensor(sigma8,dtype=tf.float32)
    init_stages = np.linspace(a0, af, n_steps, endpoint=True)
    initial_conditions = flowpm.linear_field(nc,    
                                            Boxsize, 
                                             lambda k: tf.cast(linear_matter_power(cosmo, k), tf.complex64),         
                                             batch_size=1)
    # Sample particles
    state = flowpm.lpt_init(initial_conditions, 0.1)   
    # Evolve particles down to z=0
    final_state = flowpm.nbody(state, init_stages, nc)         
    # Retrieve final density field
    state, lps_a, lps=lightcone(final_state, a[::-1], 
                                  nc, 
                                    field*60/plane_size, plane_size,
              cosmology)
    k_map=Born(lps_a,lps,ds,nc,Boxsize,plane_size,field,cosmology)
    k_map=tf.cast(k_map,dtype=tf.complex64)
    ell, power_spectrum=measure_power_spectrum_tf(k_map,field,plane_size)
    return ell, power_spectrum, k_map

# Error arises when executing this function
ell3, power_spectrum3, kmap=power_spectrum_for_cosmology(0.3075,0.8159)

A longer version of this example can be found in this notebook

Any idea of what could be wrong ?

[dlanzieri]

Hi @EiffL! I found the solution by taking your suggestion and defining the variables:

a = afactor(r)
a_s=z2a(1.00) ds=chifactor(a_s)

Inside the power_spectrum_for_cosmology function. Thank you! But I have a new problem, trying to build the function to compute the gradient (or the Jacobian), I'm running into the following error:

TypeError: Expected int32, got 0.0 of type 'float' instead.

The whole message (too long) is here notebooks/spectrum_for_cosmology.ipynb.

How can I fix it?

[EiffL]

ok, I found the problem ^^' https://github.com/modichirag/flowpm/commit/e988109bf12c54102586f27df50bf793cbe753f4#diff-48a2fc2fb081577108e9128e54c1ca8c825809137720e251080109ded0c9aea9L11

[EiffL]

I can explain later why it happened

[EiffL]

ok @dlanzieri , I think this is all solved for now: https://github.com/LSSTDESC/DifferentiableHOS/blob/main/notebooks/spectrum_for_cosmology2.ipynb

Sorry it took me a while to get to it, there were a few remaining problems, the biggest one I think is that we were computing the linear power spectrum more than necessary, and that was making things very complicated. Here I'm using instead an interpolation scheme to evaluate the linear power spectrm at the positions ndeeded for initializing the linear field.

This notebook is still approximate though, I reduced the number of steps and lensplanes to have something that went faster.

Also, loooking at our flowpm code, I found a bunch of problems in how we were handling the cosmology, and I've opened a PR to fix these issues: https://github.com/modichirag/flowpm/pull/67

I propose we follow the following strategy:

1. Merge https://github.com/modichirag/flowpm/pull/67 into FlowPM master
2. Rebase the lensing branch on the new master

[EiffL]

I'm going to close this one, because I think we did all the steps :-)