SpeedyWeather / SpeedyWeather.jl

Play atmospheric modelling like it's LEGO.
https://speedyweather.github.io/SpeedyWeather.jl/dev
MIT License
429 stars 26 forks source link

online particle tracking #413

Closed miniufo closed 6 months ago

miniufo commented 10 months ago

Many models have implemented the online Lagrangian particle tracking. Is it possible that this is also in the roadmap?

milankl commented 10 months ago

Funnily enough I just talked about that yesterday. It's not on my roadmap but definitely possible. You'd need to define how many and initial conditions (maybe also active or not) and then solve dx/dt = u with u from the model. For that you'd need interpolation onto the particle locations x but we have that functionality in RingGrids. It would really be about pulling the pieces together and thinking about how to output the particles.

I'd be happy to assist if you want to implement that.

milankl commented 10 months ago

To be more precise: I would start by defining a new particle type

abstract type AbstractParticle{NF} end

struct Particle{NF} <: AbstractParticle{NF}
    lat::NF
    lon::NF
    σ::NF
    active::Bool
end

Then adding a particles::Vector{Particle} to the prognostic variables. I'd then also create particle_advection::ParticleAdvection field in each model that contains user-facing parameters like number of particles and their initial conditions as well as an interpolator given the number of particles and the origin grid. On every time step then we need to use the u,v fields to interpolate onto the particle locations and update them following the same leapfrog time stepping used in the dynamical core. I would probably just output lat,lon,height each as an Nparticles vector that then gets a new time step on every output step effectively turning it into 3 Nparticles x Ntime matrices in the netcdf file. The rest should be done in post-processing (e.g. particle concentrations) depending on what user wants to do with the data.

milankl commented 6 months ago

@miniufo have a look at #465, particle tracking is now implemented for the 2D models (BarotropicModel, ShallowWaterModel)

In particular we have

What's still missing

miniufo commented 6 months ago

Great to know that! What I may need is only 2D tracking. So I really would like to try this. Are you going to make a new release including the PR? I am not so sure how to update my local version right now.

milankl commented 6 months ago

That's the plan. I just want to add a Scheduler that generally can be used to decide whether something should be executed on a given time step so that for example particle advection could happen at larger time steps.

miniufo commented 6 months ago

That's right. If the flow is smooth, reducing the time steps for tracking would save a lot of CPU. Do you have a doc for the tracking part now?

milankl commented 6 months ago

Compared to other calculations in the dynamical core it's still pretty cheap unless you ramp up the number of particles like crazy of course. How many are you thinking?

I haven't written any docs yet. But I can write up an example later that explains the interface.

miniufo commented 6 months ago

I would like to do a statistics of particle motion in a 2D turbulence but with free waves. So I may need many of them like every grid point for a particle (~10^5). Just curious about the efficiency of tracking. Is it multicore-enabled so that all particles are grouped and assigned to each core or likewise?

milankl commented 6 months ago

No, the ShallowWaterModel is single CPU only for now. We want to move towards single GPU within the next months and currently don't plan to build the 2D models to be efficient with multithreading. But at T85 you'll have some 20,000 grid points and typical speeds are 1,200 SYPD for the BarotropicModel and 600 SYPD for the ShallowWaterModel. At T127 with 40,000 grid points it's 300SYPD for the BarotropicModel and 150SYPD for the ShallowWaterModel. Do you need it faster?

I haven't checked how much 20,000-40,000 particles would slow things down, but I doubt it's much.

milankl commented 6 months ago

500 particles in the Galewsky jet, as you can see, one gets stuck on the North pole. Need to test that still. Do you have any capacity testing this too if I write the documentation on how to use it?

https://github.com/SpeedyWeather/SpeedyWeather.jl/assets/25530332/c67248cc-1a11-4373-a6d0-c6b84a40edb4

miniufo commented 6 months ago

That's pretty fast enough for me. So I may increase the resolution to an extent free waves can be resolved (not sure which is), then seed particles every 10 grid points or so.

But glad this is now possible. Is the standard 4th-order Runge-Kutta used for time stepping of particle position? I am expecting the docs on more details about tracking. I am glad to test this.

milankl commented 6 months ago

What do you mean by free waves?

I implemented Heun's method for the particle advection with a timestep multiples of the timestep of the dynamical core, current default is 8, which yields the video above. Meaning that while the dynamical core steps with 675s the particle advection is done every 8*675s = 90min at the resolution above (T85). I don't think that RK4 would give you much obviouss improvemenet and it requires more memory and a more complicated implementation. If you need higher accuracy I'd just reduce the time step from 8*dt to 4*dt for example, but that obviously makes the particle advection more expensive as twice as many interpolations are needed during the same model time.

milankl commented 6 months ago

5000 particles, they slow down a T85 shallow water simulation by 20% roughly

image

miniufo commented 6 months ago

That is a impressive plot (maybe I should move to Julia's ecosystem).

What do you mean by free waves?

I'm thinking the particle statistics in pure 2D turbulence and 2D shallow-water case (permits surface gravity waves). The presence of surface waves may alter the energy spectrum from classical 2D turbulence (an inverse and a forward cascades). This is a hot topic in the ocean. One may use stochastic forcing to generate various waves, or more realistically, add tidal potential forcing to generate barotropic tides.

I implemented Heun's method for the particle advection

Can you point me to more details of Heun's method? I didn't know about this.

milankl commented 6 months ago

Like the gravity wave video we have in the readme? Yeah that's possible. There's also a tutorial on it in the docs.

Heun's method is a predictor corrector method. You start with an Euler forward step with velocity at departure point and at departure time. Then you use that (preliminary aka predicted) new location to get a velocity at arrival time. Then you average both velocity at departure time and location and (predicted) arrival time and location and use that to step forward. https://en.m.wikipedia.org/wiki/Heun%27s_method

miniufo commented 6 months ago

I re-checked the readme page, the random-wave case seems fit what I need.

I see, Heun's method is the predictor-corrector method. Sometimes, using high accuracy scheme like RK4 could reduce the time step but still retain accuracy. So I am not sure which is better if computational cost and accuracy are both taken into account.

milankl commented 6 months ago

You can read the documentation on particles, their advection and tracking here

https://speedyweather.github.io/SpeedyWeather.jl/dev/particles/

(the Particle tracking section should appear in a few minutes). If you could try this, test this and generally provide any feedback that would be highly appreciated!

milankl commented 6 months ago

I didn't find the time to get my head around how to implement advection on great circles instead of local cartesian coordinates, this is described in the docs too. So instead of increasing accuracy with RK4 we should use geodetics instead. We would need to replace/update this function

https://github.com/SpeedyWeather/SpeedyWeather.jl/blob/ba429acc0e70419d423cd17a0913198a5f8d09c3/src/dynamics/particle_advection.jl#L139-L150

Any input appreciated. I don't know how expensive geodetics are, but presumably a few cos, sin evaluation are necessary. Thinking out loud, we'd need define a CartesianAdvection type and a GeodeticAdvection type and then dispatch with the particle advection using either method (I believe there's a speed/accuracy tradeoff, meaning it'll be good to have both)

miniufo commented 6 months ago

That's great! I'll try this and report here if I find something.

I guess the advection here should be OK for geodetic case, as velocities are defined as: $u = \frac{dx}{dt}=R\cos(\phi)\frac{d\lambda}{dt}$, $v = \frac{dy}{dt}=R\frac{d\phi}{dt}$