JAX MD: A Framework for Differentiable Atomistic Physics

Key Takeaways

paranoia of making sure we live because i've messed it up so many times so i'll just go to youtube just to double check i awesome so that's what i'm trying to do right now i paranoia of making sure we live because i've messed it up so many times so i'll just go to youtube just to double check i awesome i can him an echo which means fit life uh hey everybody welcome back to the jack series uh i'm super excited to host sam today sam thanks for your time and thanks for joining us i'll quickly remind everyone of what the session is introduce sam and then hand it over to him last time we met before the new year and uh we had the privilege to learn from jonathan uh this has been an ongoing course where we learn about chuck's i get to learn from incredible people like jonathan like sam and this is being hosted in collaboration with the jacks global meetup series um i wanted to quickly introduce sam and then hand it over to him sam is a senior research scientist at google he's the co-developer i believe core co-developer refinement of jack's md he has a phd in physics and you can find him on twitter at s shawlin halls uh i'm not i was telling some of the call i'm not an expert in physics at all i don't understand any details of it so he's the two expert and i won't try to make any comment but i invite everyone and i want to remind everyone please keep the questions coming in this is an interactive session and i'll be passing your questions on to sam uh sam thanks again for your time and thanks for joining us awesome thank you so much for having me um yeah so i i you know i wasn't sure um wait let me let me just get some also let me know sometimes when i i'm gonna try to do a demo uh during the talk and sometimes um when i switch slides it uh it doesn't it like loses the um the share screen so let me know if that happens anyway uh but thanks for having me um it's great to be here i've seen you know some of the talks in the series and they've i you know i i obviously love jax uh so i'm i'm really happy to be here and part of it um you know i wasn't sure what the audience was going to be like especially because you know when i give talks to more machine learning folks i think probably people don't have a lot of background in physics or in molecular dynamics so i'll throughout the talk be trying to just kind of make it a fun talk with some demos um and kind of give a gist of the kinds of work that um molecular dynamics and physics tries to like do and solve um but anyway throughout you know i uh please please ask questions i'd rather get through less stuff um and uh and have people understand than then you know get through everything and have everyone be lost so anyway uh jax md is a project we started working on maybe about three years ago and it's it's a when when jax came about i was really one of the things that i was really excited about you know because i've been using tensorflow and um i found it great for machine learning but maybe not so great if you wanted to do other stuff like science um and one of the things that i loved about jax being built on numpy was that it really seemed like an opportunity to write excellent scientific computing tools and since my background is in physics and in particular i did a lot of like atomistic physics simulations i decided to try and see how it worked um with molecular dynamics which is a kind of atomistic physics simulation um and i should say like the at the beginning it was me and my my colleague dosh working on it um but ever since then we've had you know a growing set of open source contributors so you know if you if you go to the the github and you look at the contributors i think we've now had a lot of people helping out um and doing some great work so so what is what is molecular dynamics um molecular molecular dynamics is like a cornerstone of um physics research especially like adam uh you know if you want to understand like a material or like sometimes like a protein or uh in bio in biophysics what you'll do is you'll you'll you know you know sort of like what the atoms are in the system um and you'll run a computer simulation instead of trying to do an experiment and so you could do an experiment where you like actually try to produce the material or the the protein and measure like a binding or the material and measure like you know let's say like a a glass and measure like how much does it shatter if it gets impacted by something um so you might want to run these things in experiment but that's hard because you have to like actually produce the experiment produce the material but another approach that's kind of complementary is to simulate the system in a computer and what you do is you you have all the atoms that you kind of know exists in the material or the the the system and you posit some interactions between them and usually these will be dictated by quantum mechanics which is hard to simulate so you'll use something more like approximate but if the fidelity of the of of the interactions or the the atoms is high enough you can get simulations that like agree very well with experiment and what's nice about that is if you have a simulation then you know exactly what everything is doing all the time so so you can make measurements in computer simulations that you just can't make uh in in experiments uh so this this is just an example that i ran uh in in jax md of silicon dioxide at room temperature so this would be like the glass in like a window that you might have for example and so so this is just being simulated uh this was simulated in a collab notebook um but the idea would be that if we know what all the atoms are doing we might be able to ask questions about like how hard is the is the window uh how much force might it take to shatter and if you're building something like an iphone screen or something this can be like a useful thing to do um so more often than not uh you know one of the things about physics simulations is that they are sort of stuck in the past a little bit i think um and they're built in much the same way as code was written before libraries like tensorflow or pytorch or jaxx and so what you have is you have a very fast backend that's usually written in c if you're running on cpu or cuda if you're running on gpu um and this is great because it's like super super fast and performant and they've really optimized it for the workload that they're running on but it has some disadvantages because you have significant code duplication between the two paths um and then simulations are written in a front-end language and this will often be python but it could be something else so for example there's a popular library called lamps which uses something called lamp script or its own scripting language but but in general there's some higher level front-end language that the user interacts with and again this is great because you can tailor the front end to very to to be make it really easy to run certain workloads but the problem is if you are doing something that's quite different from what these systems were designed for it's easy to fall off usability cliffs and then if you want to change them or like do something a bit different you you're forced to go down into this like c in cuda and it's it's worse right because because the code is duplicated if i write if i want to add some functionality i have to implement it in principle in both c and cuda so i have to understand this like very kind of complicated text stack um moreover all derivatives usually are handwritten uh so so you know i i don't think uh the other two uh points here i think there are positives too at this point i'm pretty convinced that we should use automatic differentiation for everything so anyway this is like i think pretty laborious and error-prone and and altogether what this means is that for a graduate student starting out in physics if they want to do something a bit different from the usual new ideas can take months uh because you have to compute all the derivatives which can be complicated you and like papers have been written about how to basically apply the chain rule in different situations and then implementing these ideas can be hard and i think it's an especially bad fit if what you want to do is combine machine learning and physics because machine learning involves complicated derivatives and simulations that maybe don't look exactly like the simulations we're used to so it's it can be somewhat difficult to merge these not impossible people do it but i think it adds a lot of friction so so what i wanted to do uh you know with this project was to take inspiration from machine learning because i think machine learning has great tools like i think the machine learning libraries overall are all fantastic compared to what people are using in the other sciences and i think there are sort of two things that have really driven these tools to be so good the first is automatic differentiation and the second is this like just-in-time compilation or like you know compiling to like a graph and then executing the graph and and having these two things be done very efficiently so with jax md i wanted to target certain high-level features i wanted it to be easy to express complicated simulations so there were two corollaries to this first simulation code should be written in the same language as experiments so if you're a researcher and you are trying to do some exotic simulation you should be able to write code that's at the same level as the rest of the library without ever leaving python and the other idea is that we wanted it to be functional so we wanted to borrow from jax because i really liked the functional style of jacks so we wanted to be functional with loosely connected primitives that can be composed and so this lets you build new simulations by composing um these sort of like primitives that are just generically useful i also we also wanted uh jacksonview to be fast enough to do high quality research as i said i think existing libraries are just super super well tuned so it's unlikely that we'll be faster than those but i think we can be sort of competitive enough that a for a researcher the the increased iteration speed of like writing research code is worth it so we do this in two ways we use you know just in time compile compilation uh to cpu gpu or tpu via xla and then as i'll talk about a little later we've implemented certain uh spatial partitioning strategies uh that let us do calculate do simulations in like order n where n is the number of atoms in the system rather than n squared uh so i'll talk don't worry about that too much i'll talk about it more uh in a little while and then finally we wanted machine learning to be a first class citizen and to us that meant two things one any function should be a neural network so if you want then a simulation to be driven by a neural network that should be fine but moreover we wanted the whole simulations to be end to end differentiable so if you wanted to do sort of like meta optimization or something where like let's say i'm simulating a material like um like glass or metal and i are a metal alloy for example and i wanted to like tune the amount of or like the structure to make it extra hard i should be able to be able to do that by differentiating through the entire simulation end hand so so these were the the goals um i i do think there are some limitations right now as i said we're slower than specialized software and not all jit back ends are created equal so the performance is not consistent across back ends i i think actually though we've had some real improvement recently and we're like maybe 30 percent or or so i think slower uh on gpu which is our best back end cpu is a lot slower um and then we're missing some simulation environments so we're always trying to increase the number of features as users want them so we're still lacking you know long-range interactions and systems with internal degrees of freedom don't worry too much about that but anyway we're always trying to add more add more so i i know this is the jack's uh talk series so probably everyone is familiar with jax but just in case there were people here who like were not familiar with jax i wanted to give a two second high level overview uh you know jax is a new machine learning library out of google it's the spiritual successor to autograd if you're familiar with that and a lot of the same folks worked on both it's basically numpy on gpu with certain composable function transformations like the gradient just in time compilation or automatic vectorization and it's and it's functional um i'll just go through like three examples of the three major function transformations grad where like i take uh a function f of x which computes let's say the dot product between two vectors x and itself and i compute df as grad of f will give me a new function that computes the gradient and so then i can either apply f of x or i can apply df of x and get the gradient um jit compiles to device so let's say i have this it returns a new function that compiles to device so let's say i have a function f of x and let's say this computes the dot product between x and itself and then it adds the dot product of x in itself so this you know this would be a stupid way to write the code but i could do it um and if i were to just execute f of x this would be executed as three separate calls to the gpu so this would be like opt-by-op mode in pi torch or eager tensorflow if i wrap f in a jit then it executes as a single call and moreover xla will optimize the code so that it won't even calculate the dot product twice probably it'll just do like twice the dot product so it'll it will do like common subs expression elimination and and get rid of the extra dot product call and finally automatic vectorization which is my personal favorite i think of the three lets you take a function that applies to a single input and apply it over a batch of inputs so here we have our function computing the dot product between a vector and itself and if we v-map f we can apply it to a matrix of like vectors or like uh an array of vectors and we'll get the dot product of the array uh with itself or each vector with itself and then this is nice because you don't need to think about batch dimensions you can write code and we use we use vmap all over the place uh in in jacksonville and i can maybe explain where we use it if it would be interesting to people because i think it's a super powerful tool okay so i guess i'd like to give a very high level overview of the design again i think this will make more sense we'll go into a demo in a second and kind of walk through the different pieces but i wanted to have like a bird's eye view before we do so so at the lowest level of every simulation we have sort of a notion of space which is what's the space in which the simulation is being performed and this might be just like euclidean space so this might just be like all the points are in two or three dimensions but often what you do in physics is you have periodic boundary conditions and you can think of that sort of like i don't know if you've played the old video game asteroids you know how you like you could go around the side of the screen and you'd wrap to the other side of the screen so often in physics because materials in the real world are huge and they have like 10 to the 23 or 10 to the 30 atoms or whatever um and we can't simulate that what we often do is we take our system and we wrap around the sides to try to like make it appear bigger than it is and so this is what's called periodic boundary conditions so so you might have like normal space you might have like periodic boundary conditions um and so so spaces is like the first thing you need to specify when you're thinking about a simulation on top of that as i mentioned we have the these these spatial partitioning systems which basically take a system of particles or atoms and separate it out into local neighborhoods and i think maybe of all the pieces of jax md this is the one that maybe is most generically useful to people for example there's like a lot of pure machine learning research that has nothing to do with molecular dynamics that is like point cloud research and i think one of the things that one of the things that you might want to do in point cloud research is separate points into local neighborhoods so you don't need to consider all pairs of points you can consider just like local regions of points so so spatial partitioning does that and i'll show you some examples later and it lets us scale to really large systems um we have vectorization utilities that are all built on top of vmap and this lets you take a um define a quantity for like one pair of atoms or one pair of particles and vectorize it over a whole system of particles and then we have neural network primitives that you can use to drive simulations uh fi then we have a bunch of different energy energy functions and i'm going to talk about this a little later but i'm going to use the term energy because that's the common term in physics if you're in machine learning uh and you come from a machine learning background energy is basically loss they're very they're very similar concepts so in in physics you define like the energy function in machine learning you'll define like the loss function um so so just keeping that in mind we have a bunch of different energy functions or a bunch of different loss functions and then then we have like higher level tools built on top of these we have a bunch of minimizers that like you know you can think of stochastic gradient descent or like momentum optimizer physical systems uh because have a bun the the minimizers that are good in machine learning are typically not very good in in physics and vice versa in physics you're always doing like full batch minimization you don't have a notion of mini batches um and you don't really have stochasticity from mini batches so typically different minimizers are better so we have a couple different minimizers then we have a bunch of simulation environments to simulate systems in different settings um and i'll talk about some of those in a little while and then we have some higher level tools to analyze the properties of these systems and i think next i have a demo if there are any questions at this point i'm happy to answer them if not i'm also happy to just dive dive right in hey sam i think i think we can get one of the questions uh surmakash he says like what are your favorite features from dx from jax um so absolutely love the map um i think that v-map so so one of the great things i think about vmap is that it lets you define a function once and then use it in so many different ways so for example when we define our notion of a space and i'll show this in a little while we define a function that computes basically the displacement vector between two points so if i have two points it calculates the vector between those points and that's like the fundamental thing that we define when we want to actually use that we use that in all sorts of ways so for example if i'm computing what um without spatial partitioning i might want to compute the displacement vector between all pairs of points in a system so that's equivalent to two to double the map if i'm computing the displacement if i'm using spatial partitioning we actually include two different formats for the spatial partitioning and the only reason why we can do that is because the two different formats for spatial partitioning correspond to two different v maps or two different ways of applying vmap to the displacement functions so like we're able to like get a lot of usage out of the same functions by by applying vmap in lots of different ways so i think that that's my favorite part i think and then i would say just um in general i really like how unopinionated jax is i guess like there's been some conversations lately about i think i i'm biased because i i think i'm biased and probably not in a way that's helpful for most people because i don't i do some machine learning but i don't do like you know my core research is not training models uh sometimes it is but i think that for example having a privileged notion of like a tf variable or like a pytorch var like i i think that's not maybe so helpful so so not like kind of some of the stuff that i really like about jax is that they've excluded things like that i also have to say i really like uh similar to vmap jax's parallelization story so like i think pmap is really pretty easy to use i haven't yet used like the pgit or xmap so i don't have opinions about those but the parallelization i think i found it very easy as like someone who's never done much parallel computing i've been able to get like a lot of large scale computing stuff working so so yeah so i think i i also yeah so i could i could talk a lot about that those are my highlights though i think have you used mask the mask transform you know i haven't um it's funny we do a lot of masking actually in in jack's md um and it's one of those things i think one of the things that there's like pluses and minuses to me having started the project or for our test having started the project so long ago uh like on the one hand it's great because we have been like maybe a little ahead of a little ahead of the curve on the other hand we've sort of had to implement a lot of stuff that the jax team maybe has implemented better afterwards so like i think we do a lot of our own masking um jacks now i i don't know what the status of the max masking is because i i like i said i haven't used it but i think like for example that's that's one place where uh maybe their solution would be a lot nicer than ours so yes it's on my list i haven't used it but i saw some spoilers so i know that you guys do really awesome stuff with that with masking yeah i i have not yet but i i do think that it's one of the features that i'm really excited uh to try out there's there's a bunch there's i i think um i i am really excited about xmap also i think it seems super super cool uh and i i think you know the the one yeah i think overall very happy i think there's like also spar like the sparse library stuff is coming along and i haven't really played around with that they've been adding so much stuff and i just haven't had time um cool yeah if there are any other questions happy happy to answer um otherwise also happy to do the demo or go start start the demo there's a very controversial question here have you looked into the julia side of what people are doing around this yeah so people i i think the most common question i get actually is is about julia um and sadly i will say that julia has been on the list of things that i really want to learn for a long time and i just like haven't i think julia seems awesome just to say i think there's certain areas where julia seems a lot better like for or have more stuff like i think for physical systems they're like i've been really impressed by their um the the things they've been working on for like different auto death um and i know there is a molecular dynamics library in julia i just haven't used it um i will say i think i think that it's a trade-off i think julia as far as i can tell is lower level it's trying to be somewhere between like c and python with a lot of nice features python is higher level and i think what that means is that there are things you can express in julia that you can't express efficiently in in python or in in jax however i think that the code you write in jaxx is very easy to understand typically and very high level and i think what that means is you lose some expressivity but what you gain is that like you you know it's really easy for someone to come along and look at the jax md code and understand sort of for the most part what's going on um i think i imagine but i'm not sure that some of the julia stuff to be efficient you need to get slightly lower level like for example i know that julia has like gpu like one of the things that i don't i don't understand again i i might be saying stuff that's wrong uh but this is my like loose impression however i will say julia seems super nice and i really wanna it's it's high on my list of things to learn i just think it's a trade-off i think you can express more stuff in julia but it's a bit lower level and i think maybe at some point i will learn from julia and then be able to say smarter stuff about it seems good okay um so yeah so so maybe i'll start doing some demo um work stuff and again i want for the most part um okay uh i wanted this demo because i wasn't sure what the background of people would be i wanted this demo to be very i'm sorry i don't see anything you don't okay great thanks for letting me know i think i need to stop and then reshare this happens sometimes when i leave keynote it um okay is that better you can see it no yeah thank you yeah so i wanted this demo to be very like kind of fun uh and not much physics uh because i thought you know i i don't know i think it's more of a machine learning crowd um and anyway i think there's nothing wrong i think one thing people think that oh if you're doing physics it has to be like very sterile but you can i think they're kind of fun okay so so what we're gonna do in the demo is uh we're gonna we're gonna try to like make a sandcastle and then blow it up that's that's the goal okay um and so what i'm gonna do first is i have this little helper function uh and this takes an image and by the way this is this notebook i put online so if you go into jackson d notebooks sandcastle uh you'll find basically this code so so the first thing we're gonna do is this this code takes an image and it fills in the pixels with like particles that i'm going to call grains of sand so it's going to return three things it's going to return the size of the image the positions of the particles or the grains and the color of the grains because they're going to correspond to the image and so one of the things that jax md has which i think might be useful for some people is it has a renderer that you can use from collab so you can just say we're going to render a an image of the size box which is our simulation volume and we're just going to put some disks with the given color at the right positions so this is this is the image that we loaded uh i will say so this this 24 here this number 24 is like how much we're discretizing the image so we took a high resolution image and we sort of sub sampled it down into these like grains some just to clarify it it's taken from an image but they are particles the image is just an image uh the image i think i can just like look at it right or can i i don't i don't know i don't know how to do this in i don't know how oh there we go oops oh what am i doing uh oh it opened to the right oh right okay so this is the image um and what i'm doing is i'm basically going into the image and i'm taking 24 by 24 blocks and i'm replacing it by a particle or a grain like so now i'm like discretizing it into particles and each particle has the color of like the pixels that it's on top okay so it's it's it is it is um so we're we're doing a discretization and and what we'll do later is we'll change the discretization so so one thing we can do is we can be like how many how many particles are there and so positions we can just like even look at we can be like what is positions so positions uh is is is an array it's a it's a jacked array and there are and and it's a we can look at the shape too if we want and we see that it's a one thousand six hundred sixteen by two so there are a thousand six hundred and sixteen grains of sand in two dimensions um and sort of as i mentioned before we the first thing we want to do whenever we're running a simulation like this is we want to define a space in which the simulation lives so so i can uh jax md has this library called space and we can make a periodic boundary condition space and there a space is really two functions it's a displacement function and the displacement function computes displacements between particles or grains so i can give it the position of the first atom and the position of the last grain and it will give me a vector between them and it also knows how to move grains so so i can give it a position of a grain and an amount that i want to displace it by and it will return a new position uh that's just shifted 10 to the to the right is really cool so what we do is we you know we take this image and then we sub sample it and so box is just the image the size of the image basically uh after sub sampling when you pass it to periodic that's kind of the like within height of uh of the space yeah exactly so this so because we're doing periodic boundaries you have to know how big the boundary is where things will start wrapping around and so that's what box tells it box if you wanted to you could also uh have box be like an array that was like 118 by 118 for example or like something else so like you could have it be like 118 by like 200 if you wanted so so box can be a scalar in which case it's assumed that the two sides are equal or it can be um a vector of dimension the size of the um space or the dimension of the space and then it specifies the different axes separately and we have other spaces that are a little more nuanced also um so yeah so so then then as i said you know we want we want to simulate the system and the way like kind of like in machine learning where if you have a problem you need to define a loss function in physics uh we need to define an energy which is going to be like dictate and and in physics you know in nature typically tries to like lower the energy of systems um and so what we want is we want an energy so that like the low low energies correspond to like the behavior that we'd like to observe um and what we're going to do is we're going to do something that's like quite common uh in physics and it's also common in machine learning so what we're going to do is we're going to write down instead of you know we have all these grains of sand instead of writing down an energy that's like a function of all the grains we're going to actually write down an energy that's a function of just two of the grains and then the total energy is going to be the sum of all pairs so this is like n squared basically um and it's going to be a sum of all pairs of grains touch in like between like the energy between each pair and you can think of all this like very similar to like um you know writing down a loss in terms of a sum over data points basically uh so we want a model like sand that's maybe like a little bit wet and what does wet sand do so the grains of sand are very hard so so if you have two grains of sand they can't interpenetrate hardly at all um and but they stick together just a little bit because there's some water and and the water causes um an interaction that that causes grains of sand to be a little bit sticky and you know if you've ever like been on a beach and like touched wet sand you know that like kind of has some form that like sticks together uh and the other thing is that grains of sand that are far away from each other don't notice each other at all so if i have two grains of sand that are far from each other on a beach they won't they don't know that they exist so we're going to use an energy function that is very common in physics it's called the leonard jones energy and so what i'm plotting here is and and so in jax md you just write energy.leonard jones and you give it the distance so it's a function only of the distance between the grains of sand so here i'm plotting the energy as as a function of the distance and what you see is for for distances less than about one the energy is really high and that means that if two grains of sand are closer than about one from each other they'll really push each other apart very hard and then if and then at when they when for very large values of um distance so when they're far apart the energy goes to zero so so far apart grains of sand don't don't touch each other at all um however grains of sand that are kind of close to each other will have an energy that's less than zero uh this is a little maybe and so um so here is where they'd like to be so so there's a little bit of attraction because grains of sand would like in principle to be uh maybe a little bit large higher than their their radius and and there's you know maybe a little bit of water just between them so that's that's what we're going for here um is is and so so you know we use this leonard jones and this is an example of vmap by the way so so we define leonard jones to be a function of the distance and then we define this helper function leonard jones pair which computes the full energy between all pairs of atoms so we give the leonard jones pair function the displacement function which says we are going to tell it how to measure displacements and then we can just feed in the positions and we can get an energy so this is like the energy of the whole system of all the of all the grains and now now we can simulate so so what we're going to do is we're going to just define some simulation parameters so we're going to simulate for 10 000 steps and we're going to write out some information about the simulation every 50 steps um and then i you can use so we're going to use a simulation environment called longivan dynamics which you might be familiar with which is used in machine learning a lot also and we give and so all the simulation environments in jax md are written in exactly the same way as all the optimizers and jacks or all the optimizers and optics so they have an initialization function that initializes the state of the simulation and a step function which takes one step of the simulation so the launch event dynamics you give it um we're gonna give it the energy and it also needs to know how to move things around and then we're gonna give it a step size we're gonna say okay we're at zero temperature and gamma is like a damping parameter so it says how damp the dynamics are so this is just a number that i chose out of thin air because it made videos that kind of look nice um but it's you know pretty arbitrary and so then what we're going to do is we're going to initialize the state of the simulation and we just give it a random key and the positions and then we're going to jit compile the step function um so this this sets up our simulation and now we just run our simulation so we're going to run for the number of steps and then we're going to take a step every step by taking one step of our step function and we're going to have a trajectory which just records the position of the atom of the grains of sand every step or every every right every step so every 50 steps okay and so now that it's running now that it's run we can we can draw the the simulation and this will just take a moment one thing that is like the slowest part actually uh which i think okay so here's the simulation and so we can see that like everything is sticking together i don't know if you can make it out i'll like zoom in but like the the particles are like wiggling around a little bit so when you first start the simulation the the sand isn't exactly where it wants to be so it starts like jostling a little bit um and we can make the simulation a little bit faster if we want to this will be useful later on which is why i'm doing it here um you you might know so so one thing about the code that we have up here is that this for loop is written in python um and and so we are just in time compiling the steps but we're not just in time compiling multiple steps together which can be a bit faster so we're going to write a slightly different simulation function and what we're going to do here is we are going to um use so jax has these loop constructs which are compiled down to single calls to the gpu and so what we're going to do is we're going to have a simulation function that takes the sand and the trajectory and we're going to write out the positions to the trajectory and then we're going to do an inner loop that takes 50 steps and just applies the step function so this takes 50 steps writes the trajectory take 50 steps and then it returns the new state and the new trajectory and this is going to be and then we're going to just run it by doing another 4i loop over the total number of remaining steps and this this will be like a little bit faster it won't be much faster but it'll be nice because now we're compiling everything so if we want to write what we're going to do in a minute is we're going to write some custom simulation code in here and we're not we're going to be exactly as fast as the launch event dynamics that we wrote that jax md comes with because we're just in time compiling everything together so this this rendering code is going to look identical uh we we almost don't need to run it but um i will okay um okay yeah so here here we go same thing happens i think it's actually like very similar trajectory okay so and i'm just gonna because i'm on my laptop and not my desktop i am going to clear the output on both of these before we go forward uh and if you know if there are any questions happy to take some now before we move on to the next step um but also happy to just think what we're good for now uh i'll ask some questions on uh at the end okay great so so yeah so now let's like blow up the sand castle which i think is the most fun part um and what we're gonna do is we're gonna have a projectile and the projectile is just gonna be a position so we're gonna start it off at i could make that zero uh at just like a third of the way up the image uh and zero in the x direction it's also going to have a radius um a strength which is like how strongly does it repel the um the grains of sand and it's going to have a velocity because like how fast is the projectile moving and in physics when you want to add something to a simulation basically you just add something to the energy so there's like this very nice analogy between like writing simulations and or like modeling nature and adding terms to the energy um and so like there's this thing i don't know if you've looked at it but there's this equation that goes around sometimes called like the standard model which is like the model for like the universe and it's just like a lot of terms in an energy function so like you know maybe the universe is just really like this um so anyway what we're gonna do is this this term is what we already had and then this term is our energy between the sand and the projectile um the grains of sand and the projectile and and so what we're going to do is we're going to um we're going to use we want the projectile to only repel the sand so unlike the sand which has this like little bit of attraction the the the um the projectile is not gonna have any attraction so we're gonna use an energy called soft sphere uh it sounds soft but we're gonna we're gonna use a thousand for the strength so it's actually not so soft and so we can plot so the blue is the leonard jones that we're using for the sand and the orange is the is the energy that we're using for the the projectile and so what you see is like basically it doesn't do anything but then if it hits the sand it pushes it away so the energy cost for having the grains of sand close to each other is very high or the grains of sand near the projectile is very high but otherwise it doesn't affect them at all and we're going to write that down uh by writing a projectile energy which gets the sand the position of the sand and the position of the projectile we're going to compute the distance between the sand and the projectile and then we're going to compute the energy based on the distance and we're going to give it the strength and we're going to give it the size of the projectile so we give it the size of the projectile we give it the strength of the interaction and we give it the distance between the two things the the sand and the projectile and then we sum over all the different grains of sand and then our total energy just as we wrote this above the total energy is the energy of the sand between the grains of sand and the energy between the sand and the projectile um and then what we're going to do is because we now have a sand and a projectile we're going to have a little bit of a state i'm just going to add this little convenience function or a class to store the state of the system so we're going to have the sand which is going to be simulated using this like nvt longivan dynamics and then the projectile we don't need anything we're just going to have it be a position um and now we're going to reinitialize our longivan dynamics but we're going to give it the total energy instead of just the energy between of the sand um and then we're going to run our simulation and this code looks almost identical to the code above but what we're going to do is our trajectory is also going to be a sandcastle and it's going to store this stand position and the projectile position at each time step and then our step function that we're going to use to actually simulate the dynamics is going to take one step of dynamics on the sand and we're going to pass in the projectile position and then we're just going to update the projectile by by the velocity so we're going to like move the projectile a little bit by the the velocity and we're going to update the step using the launch event dynamics and we're going to have a 4i loop that runs this for 50 steps and then the outer loop is going to run for the total number of simulation steps so we're going to run this simulation where does velocity come from here oh we just put it in uh up here so we just specified it as a global like just oh okay so i just chose this to look to look nice there's no it's just going to move linearly and yeah it's just going to move a little bit every every frame and you know to be honest if we wanted to make it better we would have like dt this time step be a parameter and it would really be like velocity times dt but you know just for the purposes of the demo i think uh it was okay so anyway yeah we ran we ran the simulation and now let's let's see what happens so we can we can now draw it so so the renderer can take multiple things so now we'll we'll draw the sand and we'll draw the projectile there's two separate things that's awesome why doesn't the ball reappear uh so the ball yeah great question so the ball uh doesn't reappear because here we're just doing state dot projectile plus velocity we could do shift fun of state projectile comma velocity and now if we ran this in a second it's so cool to visually see this happening yeah so i i think that's the thing i think a lot of yeah so now it'll go back around oh i feel like i'll just keep running this demo for like 5-10 hours straight just to enjoy it just to play around with velocities i have to say like i think the most time i spent like preparing the talk was like playing around with the visuals yeah um and i think that's one thing like when you're doing machine molecular dynamics normally like it's very disconnected from the visualization so i we did try quite hard to like make it feel more interactive um have a couple of questions okay the first one is is is this like visualization library can you run it outside of cola because it's super nice um that's a good question i have not made it work outside of collab i tot i think okay so the way it works is there's a javascript code that runs i mean so it will work in any ipython notebook i guess is one thing i would say um so if you were using it locally with ipython it would work there um i think it wouldn't be too much work in principle to get it to work elsewhere the problem is that it so it uses a javascript front end and then it transfers the data actually from the kernel to the front end so that part i think would need to be replaced because if you were running it externally uh you wouldn't have a kernel um yeah so so it's like a an ipython widget it's it it's yeah exactly it's an ipython widget exactly see um and so that's why when there's a lag actually the longest part is um is running is doing the data transfer actually ah that makes sense but it's super nice because i i think in my jack's 27 days of jax i did something with physics i saw it i liked it a lot and and i mean if i had this it would have been like five times more beautiful so yeah so so i guess like one thing i would say is like you know i think jacksonville is a pret i don't know pretty lightweight dependency so like people should definitely feel free to like use one part of it and discard the rest if it's not useful um so i do think that this kind of like by the way all of this also works in i have i don't know if we'll get to it but there's a it all works in 3d too um so you can do 3d that's awesome oh wow maybe i should upgrade just for fun try to upgrade what i did and see if it works out that would be awesome i'll be awesome um i i do have some there is a question i have a related question maybe of of mixing the two um like i don't know i i think well there is like it's exactly that one there's kind of different ways to do simulations like ods or finite elements um and can you give us and this one i guess it's energy based i don't know if you could call it that can you give it like an overview of when would you use each type so so there's a bunch of different simulation environments that we include and usually so i think you know one of the things in physics simulations that's it's it's kind of a subtle point but when you want to simulate a physical system you want to make sure there's sort of like a distribution that you want to be sampling from in the same way as like you know if you have like a va or something and you want to sample from the posterior there's there's a specific distribution that you would hopefully like to sample from so one of the things that we've tried hard to get right is to make sure that the distribution you're sampling from is the distribution you think you're sampling from and so that's why in these systems you typically don't just want to use a standard ode solver you can so there's actually a group that's using jax md like using the energy part of jax md and the spaces but is just using a standard like a um a runge-kutta ode solver um on top of it so so so you know it's one of those things where again like really like you can combine there are ode solvers in jax right and and like you can just use um the other parts of jackson b with those solvers the simulation environments in jax md are typically designed to be quite careful in simulating the right distribution um and so for example this system we use like this nvt longivan simulation and longivan dynamics are typically for systems where you have some uh like particles in like a solution or something so here i was like okay we have sand and it's sort of like in the air and maybe the air here is like a little too viscous but basically the idea was we want something see how like these particles sort of like slow down as they move we want something to stop everything from just like moving indefinitely we want like we we want some some friction um and so that's why i chose that here often my favorite is something called nvt nose nose hoover nvt by the way uh i should have said this earlier basically means uh that it's being simulated at a certain temperature um whereas so so he so if you just do a normal ode solver you'll be simulating at constant energy which is fine but you know if you're talking about like a system in the real world usually that's not at constant energy it's a constant temperature and so one of the things these simulation environments do is keep the system at a fixed temperature um i don't know if that helps is that like a kind of an answer i don't know um i mean i i think it does i have like an additional question like can you mix let's say an energy with an ode like kind of okay i have this od but then i'm gonna add like an energy term is that kind of a thing yeah so so one of the things you can do right is under the hood when you pass an energy to i can show you when you pass an energy what we're really doing so uh we have this library called quantity which has some useful quantities and we could be like what we do is we convert it to a force which is basically the great negative gradient so so if you did this right if i did this and then i looked at like the shape projectile okay um so so here this is now a 16 by 16 16 by two vector and this you could plug directly into an ode solver as the update oh it's really interesting but you you and what you showed us you don't you don't do this or do you we do so so jax md handles that uh for you one thing we could do for example just to show you is if i were like quantity.force total energy uh here this would be fine this would just work uh and it's just what we do is we detect internally whether you passed an energy or a force and if you pass an energy we just take the force for you because it's so common to pass energies but but yeah if you were using like a the jax ode solver for example you would just do this yourself yeah i i did that myself so i was wondering like okay and for i i did it in gravity right i was wondering okay how can i frame reframe it such that i could do the gravity simulation with uh the xmd um yeah it's possible oh yeah so if you wanted to do so like for example right we could be like gravity energy uh let's say i had a position like positions or something right uh we would want to do something like return uh