►
Description
NERSC Data Seminars: https://github.com/NERSC/data-seminars
Abstract: The upcoming generation of cosmological surveys such as DESI or LSST will probe the Universe on an unprecedented scale and with unparalleled precision, to answer fundamental questions about Dark Matter and Dark Energy. However, optimally extracting cosmological information from this massive amount of data remains a major challenge, and constitutes a very active research area. Having access to differentiable forward simulations of these surveys paves the way to novel and extremely powerful gradient-based inference techniques. For instance, we have demonstrated potential for over a 50% information gain in constraining Dark Energy using the upcoming DESI galaxy survey. In this talk, we will present FlowPM, the first differentiable cosmological N-body simulation code implemented in TensorFlow for seamless integration with deep learning components and gradient-based inference techniques. After showcasing a few examples of the benefits of such a tool, we will discuss our efforts to scale these simulations to large supercomputers using the Mesh TensorFlow framework.
A
A
B
B
D
B
B
So,
let's
start
with
the
basic
change:
this
is
the
so
in
this
morning,
what
we
want
to
do
is
essentially
try
to
study
the
birth
and
evolution
of
the
universe'.
We
know
the
big
picture
and
principle
that
there
was
a
big
bang
which
basically
gave
rise
to
the
universe,
followed
by
an
innovation.
We
have
the
universe,
expanded
and
then,
after
that,
essentially
the
IND
expanding
universe.
The
dark
matter
not
collapsed
to
form
a
doze
and
galaxies
that
we
observe
so
based
on
these
things.
B
B
How
we
approach
this
in
cosmology
is
basically
four
different
kind
of
cosmological
terms.
First
of
them
is
the
cosmic
microwave
background
radiation.
This
is
essentially
the
first
light
of
the
universe,
which
was
generated
by
the
very
thinking
of
the
universe
from
there.
So
this
is
the
for
this
that
you
can
see
in.
B
Than
that
we
have
other
survey,
so
you
have
like
galaxy
clustering
surveys.
These
are
big
telescopes
that
are
essentially
mapping
out
the
universe
to
see
where
how
the
galaxies
are
clustered
together.
Second,
you
can
have
supernovas,
which
are
basically
stars
going
boom,
so
these
are
explosions
which
outshine
the
entire
galaxies
first
short
periods
of
time,
and
they
allow
us
to
see
in
the
local
regions
to
the
farthest
reaches
of
the
universe
and
allows
us
to
constrain
the
expansion
of
the
universe.
Next,
we
have
gravitational
lensing.
B
This
is
essentially
bending
of
the
light
due
to
the
dark
matter,
as
it
comes
from
the
distance
galaxies
to
us
and
with
this
bending
of
the
light,
this
is
one
of
the
only
probes
that
you
can
essentially
probe
the
Dark
Matter
directly,
but
other
else
for
everything
else.
We
essentially
rely
on
the
for
transport
here
on
the
visible
matter.
What
this
gives?
You
is
kind
of
the
dark
Hydra
directly,
so
where
we
are
going
with
this
is
essentially
we
have.
B
B
So
this
is
a
14,000
square
degree
surveys
over
five
years,
14,000
degrees
as
square
degrees,
is
essentially
one
third
of
the
total
that
it
would
be
mopping
up
and
a
cube
first
light
last
one.
So
this
is
the
first
picture
of
the
sky
that
it
took
and
the
first
spectrum
this
right.
So
in
five
years
it
will
map
thirty
five
million
galaxies.
We
will
compare
the
previous
largest
spectroscopic
survey
that
we
had
boss,
which
concluded
a
couple
years
back
today,
is
about
1.5
million
galaxies.
This
is
in
just
one
year.
B
C
B
Be
this
is
spinning
90.
The
blue
curve
here
is
the
best
constraints
we
have
from
the
previous
survey.
The
red
curve
is
what
we
get
from
desi
at
the
end
of
five
years.
The
point
at
the
center,
like
one
comma,
zero
point,
minus
one
comma
zero
point
is
for
the
lack
of
better
word,
not
interesting
science,
but
everything
other
than
that.
If
we
have
constraints
other
than
that,
then
it's
a
meaning
science
that
we
do
not
know.
Yet.
B
This
is
going
to
be
the
largest
survey
that
we
have
ever
done
is
going
to
get
thousand
images
each
night,
which
will
generate
15
terabytes
of
data
for
10
years,
15
terrible
day
for
night
for
10
years
between
thousand
square
degrees,
and
it
will
observe
the
entire
sky
once
every
few
days,
which
will
essentially
give
you
ten
billions
of
objects,
each
observe
the
thousand
times.
So
the
point
I'm
trying
to
make
here
that
the
precision
for
this
is
going
to
be
huge,
and
given
that
this
will
generate
so
much
data,
you.
E
B
Is
again
the
same
energy
constraints?
Rent
is
the
curve
radius,
the
constraints
that
we
have
here
and
the
black
is.
What
is
the
potential
of
these
surveys
to
give
us
in
the
future
right?
So
the
modern
cosmological
surveys
will
have
unprecedented
istic,
allow
our
that's
without
saying,
and
we
have
great
potential
for
new
discoveries,
so
data
is
no
longer
a
problem.
They're
like
we
have
enough
data
at
this
point,
we'll
get
an
update
at
this
point.
The
problem
is
how
to
best
use
it
and
Mac
brings
in
a
lot
of
challenges.
B
So,
just
to
briefly
outline
some
of
the
challenges
that
we'll
be
facing
is
we
have
different
cosmological
probes
that
will
be
tracing
the
universe
and
you
want
to
combine
them
together
because
you
brake
degenerates,
you
become
robust
against
systematics
and
stuff,
so
you
want
to
combine
them
together,
but
since
they
are
tracing
the
same
universe,
they
are
correlated
to
each
other,
and
that
brings
a
problem.
So
you
have
to
come
when
you're
doing
inference.
You
want
to
calculate
the
joint
covariance
which
basically
takes
all
the
parameters
into
account
and
that's
a
huge
problem.
B
So
that's
what
I'm
trying
here
this
is
a
sample
covariance
matrix,
which
are
solid
lines,
and
you
see
that
there
are
huge
of
Daniel
components.
The
way
you
did
mostly
get
covariance,
matrices
and
cosmology
is
by
running
simulations.
But
here
now
you
have
huge
challenges,
because
you
want
to
get
precise.
Simulations
will
generate
all
those
five
observables
at
the
same
time
with
enough
precision.
B
So
the
next
question
that
you
have
is
like
what
are
the
optimal
summary
statistics
that
you
want
to
use
to
get
the
information
from
this
data.
So
this
is
showing
slice
of
an
n
body
simulation
at
the
present
time.
And
if
you
look
at
this
field,
this
field
is
anything
but
abortion
field
and
the
most
favorite
statistic
that
cosmologists
use
this
two
point
function
or
the
power
spectrum
of
anything
which
is
an
optimal
summary
statistic.
If
you
have
a
version
field,
but
this.
B
But
that
so
you
can
go
to
higher
order
functions.
You
can
go
to
three-point
functions,
you
can
go
to
eat
pants
and
other
kind
of
things,
but
again
you
do
not
know
which
of
them
is
optimal
and
secondly,
how
do
you
combine
for
them
because
again
they
are
getting
some
of
the
information
is
all
of
them
and
then
the
problem
is:
how
do
you
do
the
analysis?
Much
have
this
so
the
most
common
way
you
do.
B
It
is
like,
once
you
have
the
summary
statistic
of
your
choice:
you
right
now,
I'm
gonna
take
a
likelihood
for
that
and
then
you're
done
I'm
CMC
chain
to
get
the
posterior
on
the
parameters
that
you're
interested
in.
So
this
is
I'm
paying
an
example
of
posterior,
but
this
severely
constrained
to
you,
because
now
you
have
now
we
are
focused
on
writing
Donner
analytics,
which.
C
B
Mean
much
more
than
that,
but
there.
So
there
are
multiple
inference.
One
of
the
important
things
that
we
can
do
is
you
can
use
the
reconstruction
which
are
dictated
slightly
day
later
on.
Then
we
construct
the
initial
conditions
of
the
universe
and
you
can
look
at
the
previous
figure
with
the
CMV
that
was
then.
You
know
that
the-
and
so
this
is
one
of
the.
F
F
B
I
mean
right
now
be
by
using
depend
on
the
observable
of
interest.
We
have
different
summary
statistics
and
people
are
trying
to
combine
them.
I
will
talk
quickly
about
like
using
deep
learning
models
to
learn
that
likelihood
and
compare
them.
Fine
I
was
not
trying
to
give
the
cosmology
too
minimal
for
this,
but
yeah
we
can.
There
are
works
later,
okay,
so
one
way
to
do
that
is
to
make
a
forward
warning
approach
here.
B
What
you
do
is
that,
instead
of
putting
all
your
eggs
in
one
basket
and
trying
to
get
the
analytic
likelihood,
you
take
a
step
back
and
then
you
try
to
forward
model
the
entire
observations
that
you
have.
The
advantages
of
this
is
that
you
can
combine
different
cosmological
probes
together,
and
you
can
also
break
your
poor
model
into
different
components,
each
of
which
is
like
each
of
which
is
tractable.
So
what
I'm,
showing
here
is
an
example
to
is
an
example
for
credible
probabilistic
came
order
to
get
the
galaxy
shear
here.
These.
C
B
Can
parameterize
all
of
them
with
latent
parameters,
but
the
problem
here
becomes
as
that,
then
the
number
increases
like
a
high
dimensional
problem
for
insurance.
This
has
been
a
challenge
for
a
long
time,
but
recently
a
lot
of
methods
have
been
proposed
which
can
being
this
packable.
The
common
thing
for
all
of
them
is
that
you
need
fast,
accurate
differentiable
for
simulations,
so
this
is
one
step
where
you
can
use
them
the
next
thing
where
they
can
useful.
This
is
just
so.
B
This
is
essentially
one
of
the
examples
from
the
previous
like,
but
this
is
important
for
us
because
we
have
spent
a
lot
of
time
for
the
last
three
years.
Thinking
about
the
7b
CCP
here.
What
we
want
to
do
is
essentially,
if
you
have
a
forward
model
to
go
from
the
initial
conditions
to
any
of
the
observed
data
you
want
to
reconstruct.
You
want
to
use
that
where
you
construct
some
kind
of
latent
cosmological
fields,
which
it
will
not
see
and.
B
And
the
simplest,
where
you
would?
How
would
you
approach
what
is
your
drive,
a
posterior
personation
I'm,
starting
from
initial
conditions,
because
I
know
there
are
questions
of
the
exact
provider
on
them
and
you
write
the
likelihood
which
is
basically
in
a
compound
rate
against
the
forward
modern
simulation,
the
challenges
here.
This
is
a
high
dimensional,
optimization
problem.
B
B
B
You
find
with
of
the
people
and
over
the
last
three
years
we
have
applied
it
to
different
methodological
proxies,
weak
lensing,
21-centimeter
intensive
mapping
for
different
angles.
Different
was
not
the
point
I
am
trying
to
make
here
is
that
this
is
a
useful
is
first
and
we
have
the
framework
for
this,
but
to
go
beyond
what
we
have
right
now.
B
We
need
interfacing
with
D
plan,
so,
okay,
so
I
think
at
this
point
we
have
reached
a
point
that
I've
convinced
you,
hopefully
that
we
need
cosmological
forward
models
which
are
differentiable
so
taking
this
step
back,
what
are
small
as
you
go
forward
models.
So
what
happened?
What
are
you
doing
physiology?
Is
you
start
with
any
shape
indifference
of
the
universe,
and
then
you
try
to
evolve
the
particles
and
the
gravitational
force.
This
is
a
slice
under
is
a
movie
from
the
illustris
simulation
on
the
left.
B
B
Well,
as
you
can
see,
it's
like
whenever
you
have
a
variable
here,
but
getting
this
technique
is
super
expensive
because
you
have
to
take
into
account
all
the
hydrodynamics
or
the
other
forces
that
go
in
along
with
the
gravity,
and
these
are
super
expensive
simulations,
so
you
cannot
run
them
over
and
over
again
multiple
times.
So
what
we
do
in
cosmology
is
the
next
best
thing.
B
The
following
models:
in
this
knowledge
I
mean
you
start
with
a
linear
field.
So
this
is
the
linear,
Dark
Matter
field
here,
I
have
not
put
in
any
way
ants
at
this
point,
you
will
do
the
n-body
simulation.
They
get
the
final
battlefield
and
after
this,
depending
on
the
observable
of
interest,
your
word
models
can
diverge
most
often.
What
do
you
do?
Is
you
find
out
my
healers?
These
are
the
high
density
regions
in
the
Dark
Matter
field,
which
are
collapsed
under
gravity
because
that's
where
everything
moves.
C
B
The
distribution
approaches
are
not
differentiable,
so
that's
the
first
step
where
it
learning
canal
is
that
the
end
of
the
fabric
differentiable
forward
models
here.
So
what
I'm,
showing
here
again
is
the
initial
conditions.
The
final
data
field,
the
Halos
galaxies
in
this
case-
that
you
have
sampled.
B
C
B
B
C
B
Problem
so
when
this
works,
what
I
use
here
because
version
like
they
were
trying
to
make
a
life
here,
I
had
to
smooth
my
data
on
scales,
which
is
losing
a
lot
of
information.
So
what
do
you
want
to
do?
Next
is
essentially
go
a
step
further
and
one
is
on
the
light
layer.
So,
for
example,
you
can
write
as
d
planning
to
learn
the
likelihood
of
given.
B
Is
the
likelihood
of
a
galaxy
field
that
I
see
and
while
I'm
facing
it
in
terms
of
galaxies,
you
can
do
this
with
any
observable
that
you're
interested
in
so
one
way
you
can
do
it
is
use
a
mixture
density
Network.
What
it
does
is.
It
mainly
takes
on
an
input.
So
what
you
do
here?
Is
you
parameterize
your
likelihood
for
the
data
with
a
mixture
model
say,
for
example,
omikoshi
mixture,
and
then
you
are
trying
to
learn
the
parameters
of
that
like
little
model,
given
an
input
vector
field
right
and.
E
B
Last
function
that
you
optimize
is
essentially
the
record
of
the
data.
So
once
you
have
this,
you
can
use
this
back
for
the
in
our
case,
for
the
reconstruction
exercise,
you
can
go
a
step
further.
You
can
use
pixel
CNN.
This
allows
what
Excel
CNN
does
is
basically
write
down
the
likelihood
of
the
data
as
an
autoregressive
model,
so,
for
example,
the
probability
of
having
an
observation
at
x.
I
depends
on
all
the
points
in
the
data
that
you've
seen
before.
What
this
does
is
allows
you
to
capture
correlations
in
your
noise
Angelita.
C
B
You
do
not
have
to
stick
to
the
diagonal
like
here
that
can
be
used
so
far,
so
we
have
this
working.
This
is
just
showing
the
galaxy
positions
so
on
the
left
is
the
axis
that
I
find
with
our
traditional
gospel
is
equal
model.
The
unite
is
a
sample
that
I
generated
from
you
like.
We
would
learn
with
my
students
anymore,
so.
B
B
B
Okay,
so
that
was
the
cosmology
part
of
it.
Now
I'll
basically
discuss
what
we
have
done
so
to
go
next,
so
taking
a
step
back,
n-body
simulations
are
hard
and
expensive.
The
way
you
approach
them
in
cosmology
is
with
the
next
best
thing.
What
we
have
is
a
particle
motion
simulation
and
particle
in
my
simulation.
What
you
do
is
you
put
particles
on
a
grid
and
calculate
gravitational
forces
on
a
regular
tree.
B
You
lose
small
scales
based
on
the
waves
resolution
in
this
case,
but
on
the
large
scales
thing
is
still
are
correct,
and
so
the
numerical
scheme
for
this
is
essentially
the
set
that
you
interpolate
particles
from
the
grid.
Let's
commit
the
pass
density
and
the
grid
the
interpolation
scheme
that
people
mostly
use
as
a
clown
itself
binning.
So
this
particle
exploded
into
like
four.
It's
giving
me
with
the
four
nearest
points
in
2d.
C
B
3D,
once
you
have
the
mass
on
a
regular
grid,
you
can
use
fast.
Fourier
transforms
to
calculate
the
gravitational
forces.
So
this
is
just
the
Poisson
equation,
which
is
simulating
gravitational
potential
to
the
mass
density
and
you
take
the
Fourier
transform
to
get
V
force,
but
the
colonists
in
the
room
I
am
taking
some
liberties
here.
B
Next,
once
you
have
the
forces
you
into
putting
them
back
in
the
grid
points,
so
once
you
have,
the
forces
on
the
grid
point
change
every
time,
the
particle
positions,
and
then
you
contain
the
momentum
of
the
particle.
So
this
is
just
updating
the
momentum
in
the
force
and
tracking
function,
so
the
K
Kappa
here
is
a
scalar
and
once
you
have
updated
the
moment.
G
B
You
can
update
the
positions
of
the
particles,
so
this
is
one
time
step
in
a
particle
in
my
simulation
and
then
you
just
repeat
it,
but
how
about
you
on
so?
The
whole
point
of
this,
then,
is
that
in
simulations
is
a
series
of
operators
and
the
only
do
critical
operations
you
need
fast
Fourier
transforms
are
what
basically
gain
you
a
lot
about.
The
n-body
simulations
in
this
case,
but
when
you
are
doing
is
also.
B
C
B
B
B
So
the
difference
here
from
the
previous
code
is
that
I
made
my
linear
input
field,
a
video
which
is
trainable,
and
this
is
just
doing
the
same
n
body
simulation
to
get
the
final
feed.
Depending
on
your
likelihood.
You
just
write
down
the
residual
and
the
likelihood
that
you
have
the
prior
that
you
have
and.
C
B
B
B
B
There
are
reasons
to
believe
that
at
some
point
you
do
not
have
a
unique
solution.
So
that's
what
we
did
like
do
you
have
a
paper
or
not
like,
depending
on
the
noise
level
and
the
initial
resolution,
you
do
not
have
any
exception.
The
simplest
way
you
can
think
of
it
is
that
if
you
have
two
particles
which
have
after
a
shell
crossing,
you
do
not
know
if
they
like
did
a
crossing
or
they
came
from
or
if
they
came
from
so
on,
small
scales
will
definitely
do
not
have
uniqueness
of
the
solution.
B
Right,
so
this
is
essentially
considering
your
patient
okay,
so
because
the
challenges
that
so
many
kinds
of
interesting
sighs
do
not
fit
on
the
GPU
everything
that
I've
shown
you
so
far
that
we
have
been
working
like
that.
We
have
had
for
like
last
once
at
this
point,
I
will
be
using
that
single
GPU,
so
the
whole
reconstruction
framework
with
a
128
cubed
simulation,
which
fits
on
one
cheek.
You
are
using
hole
down
here
on
earth,
so
that
is
operational,
but
any
survey
of
interesting
size
will
not
fit
unlike
its.
B
You
have
to
go
at
least
8,000
Q,
depending
on
the
observation
like
for
some
observations
like
21
centimeter
interesting,
we
have
to
go
to
10000
simulations
and,
depending
on,
like
it's
a
higher
dimensional,
optimization
so,
depending
on
the
optimization
scheme
that
you
want
to
use.
If
you
are
using
something
like
the
current
inference
machine
you
patches
of
this
thousand
of
simulations.
So
essentially
the
point.
B
B
C
B
C
B
Tracks
so
this
is
a
super
duper
example
to
illustrate
this.
This
is
a
simple
two
layer
fully
connected
Network,
so
all
you
are
doing
is
X
time,
which
gives
you
the
hidden
layer
and
the
hidden
city,
output,
Y.
So
and
the
first
lines
here,
you
are
defining
the
dimensions
that
you
display,
your
name
in
all
the
dimensions
in
the
second
part
of
this,
you
are
getting
the
densest
with
the
dimensions
that
you
excited
and
then
the
last
part
is
the
logical
operation
that
we
would
do
to
write
down
this
fully
connected
Network.
B
F
C
B
B
B
B
C
B
They
all
require
a
please.
So
if
you
want
to
go
for
model
parallelism
on
Monday
scaling
missing
any
of
these,
this
is
a
good
choice.
To
begin
with,
so
we
started
investigating
this.
This
is
a
first
Superman
implementation
of
every
which
works.
All
we're
doing
here
is
essentially
combining
effecti
along
the
three
dimensions.
B
Where
you
do
one
FFT
along
every
dimension
once
and
then
you
reshape
and
spirit.
There
is
and
then
do
the
domain
X,
okay
from
the
next
dimension,
and
so
you
need
to
do
this
transpose
and
reshape
operation,
because
the
direction
that
you
hang
the
FFP
n
has
to
be
on
the
same
processor,
and
since
you
have
spread
them,
only
one
of
the
dimensions
can
be
like.
B
I'm
certain
that
we
can
do
better
than
this
one
thing
you
can
also
do
is
like
implement
the
transpose
and
treat
it
as
like
a
single
mesh
operation
by
being
somewhat
clever
to
be.
But
the
whole
point
is
that
this
is
a
very
different
communication
from
the
general
planning
cases,
because
here
you
are
essentially
community
communicating
the
full
tensor
of.
B
C
H
B
That
I'm,
using
is
a
symbol
of
one
day
match
I'm,
just
splitting
my
X
dimension.
Of
course,
all
those
eight
processes,
I
think
embodies
you.
They
use
the
device
they
spend
implementation
for
the
and
they
for
GP,
offset
communication,
and,
though
I'm
not
saying
here
when
you
go
from
one
node
to
to
load
things.
B
B
B
B
B
I
E
The
usual
answer,
one
point
also
to
make
here,
is
that
we
have
different
implementation
of
how
this
mesh
is
being
run
on
GPUs
and
we're
using
this
as
I
am
the
mesh
implementation,
which
is
essentially
doing
something
similar
to
an
MPI
communication,
so
between
different
nodes
or
different
devices,
but
in
the
first
implementation
were
using
some
exactly
different,
where
we
have
Joel,
perhaps
executed
by
different
ICS,
but
all
is
managed
by
a
single
process.
So
it's
not
as
scaleable
implementation
that
we're
using
on
core
GPU.
So
that's
that's.
Why
we're
interested
in
playing
with
GPUs?
C
E
B
The
first
thing
is
optimized
TV
from
the
kitchens
that
I've
know
they're
using
to
be
fourth
in
central
communication,
but
maybe
you
can
gain
have
a
comp
like
using
the
NPI,
and
second
is
the
measurement
diction.
That
transfer
just
said
if
the
device
and
implementation
is
actually
scale,
it
wasn't
for
like
exists,
KT
computing,
encoding
and
so.
C
E
This
stage
way
to
extract
something
out
of
the
computational
graph
is
to
turn
her
into
a
tensor
that
is
fully
replicated
across
all
devices.
So
it
needs
to
be
small
enough
and
then
you
can
extract
it
into
a
normal
tensor
flow.
So
that's
fine.
If
you're
looking
at
some
summary
statistics
of
your
data
and
for
us
it
might
be
enough,
but
if
you
want
to
save
your
simulation
output,
you
need
to
have.
C
G
B
D
D
C
D
E
B
C
G
B
B
B
B
B
H
H
B
C
B
He
essentially
reinventing
the
wheel
again
because,
if
you
want
to
say,
use
pixel
C
and
then
you
have
to
quote
the
derivatives
of
the
pixels.
So
that's
why
we
want
to
do
like
just
focus
on
developing
forward
models.
Well,
I
think
they,
whatever
I'd,
take
care
of
the
whole
radiant
part
of
it.
Let's
see
and
that
motivation
of
doing
it.