def result(self): grid = Grid.create(self.dimensions, self.box_size.as_vector_with_length(3)) zero = numpy.zeros(self.dimensions).flatten() | units.m / units.s rho, rhovx, rhovy, rhovz, rhoe = self.sph_code.get_hydro_state_at_point( grid.x.flatten() - self.box_offset, grid.y.flatten() - self.box_offset, grid.z.flatten() - self.box_offset, zero, zero, zero) grid.rho = rho.reshape(self.dimensions) grid.rhovx = rhovx.reshape(self.dimensions) grid.rhovy = rhovy.reshape(self.dimensions) grid.rhovz = rhovz.reshape(self.dimensions) grid.energy = rhoe.reshape(self.dimensions) if self.do_scale: grid.rho *= self.sph_code.gas_particles.total_mass() / (grid.rho.sum() * grid.cellsize().prod()) total_sph_momentum = (self.sph_code.gas_particles.mass.reshape((-1,1)) * self.sph_code.gas_particles.velocity).sum(axis=0) total_grid_momentum = grid.momentum.reshape((-1,3)).sum(axis=0) grid.momentum += total_sph_momentum / grid.cellsize().prod() - total_grid_momentum grid.energy *= self.sph_code.thermal_energy / (grid.energy.sum() * grid.cellsize().prod()) return grid
def result(self): grid = Grid.create(self.dimensions, self.box_size.as_vector_with_length(3)) grid.add_vector_attribute("momentum", ["rhovx","rhovy","rhovz"]) zero = numpy.zeros(self.dimensions).flatten() | units.m / units.s rho, rhovx, rhovy, rhovz, rhoe = self.sph_code.get_hydro_state_at_point( grid.x.flatten() - self.box_offset, grid.y.flatten() - self.box_offset, grid.z.flatten() - self.box_offset, zero, zero, zero) grid.rho = rho.reshape(self.dimensions) grid.rhovx = rhovx.reshape(self.dimensions) grid.rhovy = rhovy.reshape(self.dimensions) grid.rhovz = rhovz.reshape(self.dimensions) grid.energy = rhoe.reshape(self.dimensions) if self.do_scale: grid.rho *= self.sph_code.gas_particles.total_mass() / (grid.rho.sum() * grid.cellsize().prod()) total_sph_momentum = (self.sph_code.gas_particles.mass.reshape((-1,1)) * self.sph_code.gas_particles.velocity).sum(axis=0) total_grid_momentum = grid.momentum.reshape((-1,3)).sum(axis=0) grid.momentum += total_sph_momentum / grid.cellsize().prod() - total_grid_momentum grid.energy *= self.sph_code.thermal_energy / (grid.energy.sum() * grid.cellsize().prod()) return grid
def make_grid(number_of_grid_cells, length, constant_hydrogen_density, inner_radius, outer_radius): grid = Grid.create([number_of_grid_cells] * 3, length.as_vector_with_length(3)) grid.radius = grid.position.lengths() grid.hydrogen_density = constant_hydrogen_density grid.hydrogen_density[grid.radius <= inner_radius] = 0 | units.cm ** -3 grid.hydrogen_density[grid.radius >= outer_radius] = 0 | units.cm ** -3 return grid
def setup_simple_grid(self): test_grid = Grid.create((4,3,2), [1.0, 1.0, 1.0] | units.m) test_grid.rho = numpy.linspace(1.0, 2.0, num=24).reshape(test_grid.shape) | units.kg/units.m**3 test_grid.rhovx = test_grid.rho * (3.0 | units.m/units.s) test_grid.rhovy = test_grid.rho * (4.0 | units.m/units.s) test_grid.rhovz = test_grid.rho * (0.0 | units.m/units.s) test_grid.energy = test_grid.rho * ((1.0 | (units.m/units.s)**2) + 0.5 * (5.0 | units.m/units.s)**2) return test_grid
def hydro_plot(view, hydro_code, image_size, figname): """ view: the (physical) region to plot [xmin, xmax, ymin, ymax] hydro_code: hydrodynamics code in which the gas to be plotted is defined image_size: size of the output image in pixels (x, y) """ if not HAS_MATPLOTLIB: return shape = (image_size[0], image_size[1], 1) size = image_size[0] * image_size[1] axis_lengths = [0.0, 0.0, 0.0] | units.m axis_lengths[0] = view[1] - view[0] axis_lengths[1] = view[3] - view[2] grid = Grid.create(shape, axis_lengths) grid.x += view[0] grid.y += view[2] speed = grid.z.reshape(size) * (0 | 1/units.s) rho, rhovx, rhovy, rhovz, rhoe = hydro_code.get_hydro_state_at_point( grid.x.reshape(size), grid.y.reshape(size), grid.z.reshape(size), speed, speed, speed) min_v = 800.0 | units.km / units.s max_v = 3000.0 | units.km / units.s min_rho = 3.0e-9 | units.g / units.cm**3 max_rho = 1.0e-5 | units.g / units.cm**3 min_E = 1.0e11 | units.J / units.kg max_E = 1.0e13 | units.J / units.kg v_sqr = (rhovx**2 + rhovy**2 + rhovz**2) / rho**2 E = rhoe / rho log_v = numpy.log((v_sqr / min_v**2)) / numpy.log((max_v**2 / min_v**2)) log_rho = numpy.log((rho / min_rho)) / numpy.log((max_rho / min_rho)) log_E = numpy.log((E / min_E)) / numpy.log((max_E / min_E)) red = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum( numpy.zeros_like(rho.number), log_rho)).reshape(shape) green = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum( numpy.zeros_like(rho.number), log_v)).reshape(shape) blue = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum( numpy.zeros_like(rho.number), log_E)).reshape(shape) alpha = numpy.minimum( numpy.ones_like(log_v), numpy.maximum( numpy.zeros_like(log_v), numpy.log((rho / (10*min_rho))) ) ).reshape(shape) rgba = numpy.concatenate((red, green, blue, alpha), axis=2) pyplot.figure(figsize=(image_size[0]/100.0, image_size[1]/100.0), dpi=100) im = pyplot.figimage(rgba, origin='lower') pyplot.savefig(figname, transparent=True, dpi=100, facecolor='k', edgecolor='k') print "\nHydroplot was saved to: ", figname pyplot.close()
def hydro_plot(view, hydro_code, image_size, figname): """ view: the (physical) region to plot [xmin, xmax, ymin, ymax] hydro_code: hydrodynamics code in which the gas to be plotted is defined image_size: size of the output image in pixels (x, y) """ shape = (image_size[0], image_size[1], 1) size = image_size[0] * image_size[1] axis_lengths = [0.0, 0.0, 0.0] | units.m axis_lengths[0] = view[1] - view[0] axis_lengths[1] = view[3] - view[2] grid = Grid.create(shape, axis_lengths) grid.x += view[0] grid.y += view[2] speed = grid.z.reshape(size) * (0 | 1 / units.s) rho, rhovx, rhovy, rhovz, rhoe = \ hydro_code.get_hydro_state_at_point( grid.x.reshape(size), grid.y.reshape(size), grid.z.reshape(size), speed, speed, speed) min_v = 800.0 | units.km / units.s max_v = 3000.0 | units.km / units.s min_rho = 3.0e-4 | units.g / units.cm**3 max_rho = 0.1 | units.g / units.cm**3 min_E = 1.0e11 | units.J / units.kg max_E = 1.0e13 | units.J / units.kg v_sqr = (rhovx**2 + rhovy**2 + rhovz**2) / rho**2 E = rhoe / rho log_v = numpy.log((v_sqr / min_v**2)) / numpy.log((max_v**2 / min_v**2)) log_rho = numpy.log((rho / min_rho)) / numpy.log((max_rho / min_rho)) log_E = numpy.log((E / min_E)) / numpy.log((max_E / min_E)) red = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_rho)).reshape(shape) green = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_v)).reshape(shape) blue = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_E)).reshape(shape) alpha = numpy.minimum( numpy.ones_like(log_v), numpy.maximum(numpy.zeros_like(log_v), numpy.log( (rho / (10 * min_rho))))).reshape(shape) rgba = numpy.concatenate((red, green, blue, alpha), axis=2) pyplot.figure(figsize=(image_size[0] / 100.0, image_size[1] / 100.0), dpi=100) im = pyplot.figimage(rgba, origin='lower') pyplot.savefig(figname, transparent=True, dpi=100) print("\nHydroplot was saved to: ", figname) pyplot.close()
def setup_particles_in_nbodycode(self): staggered_grid_shape = numpy.asarray(self.gridcode.grid.shape) + 2 corner0 = self.gridcode.grid[0][0][0].position corner1 = self.gridcode.grid[-1][-1][-1].position delta = self.gridcode.grid[1][1][1].position - corner0 print delta.prod() self.volume = delta.prod() staggered_corner0 = corner0 - delta staggered_corner1 = corner1 self.staggered_grid = Grid.create( staggered_grid_shape, staggered_corner1 - staggered_corner0 + delta) # self.staggered_grid.x += staggered_corner0.x # self.staggered_grid.y += staggered_corner0.x # self.staggered_grid.z += staggered_corner0.x # self.staggered_grid.p000 = 0.0 | mass # self.staggered_grid.p100 = 0.0 | mass # self.staggered_grid.p010 = 0.0 | mass # self.staggered_grid.p110 = 0.0 | mass # self.staggered_grid.p000 = 0.0 | mass # self.staggered_grid.p100 = 0.0 | mass # self.staggered_grid.p011 = 0.0 | mass # self.staggered_grid.p111 = 0.0 | mass # self.staggered_grid.p001 = 0.0 | mass # self.staggered_grid.p101 = 0.0 | mass self.normal_grid = Grid.create( self.gridcode.grid.shape, corner1 - corner0 + delta) particles = Particles(self.normal_grid.size) particles.mass = 0.0 | mass particles.x = self.grid.x.flatten() particles.y = self.grid.y.flatten() particles.z = self.grid.z.flatten() particles.vx = 0.0 | speed particles.vy = 0.0 | speed particles.vz = 0.0 | speed particles.radius = 0.0 | length self.nbodycode.particles.add_particles(particles) self.from_model_to_nbody = particles.new_channel_to( self.nbodycode.particles) self.nbodycode.commit_particles() self.particles = particles
def setup_particles_in_nbodycode(self): staggered_grid_shape = numpy.asarray(self.gridcode.grid.shape) + 2 corner0 = self.gridcode.grid[0][0][0].position corner1 = self.gridcode.grid[-1][-1][-1].position delta = self.gridcode.grid[1][1][1].position - corner0 print delta.prod() self.volume = delta.prod() staggered_corner0 = corner0 - delta staggered_corner1 = corner1 self.staggered_grid = Grid.create( staggered_grid_shape, staggered_corner1 - staggered_corner0 + delta) #self.staggered_grid.x += staggered_corner0.x #self.staggered_grid.y += staggered_corner0.x #self.staggered_grid.z += staggered_corner0.x #self.staggered_grid.p000 = 0.0 | mass #self.staggered_grid.p100 = 0.0 | mass #self.staggered_grid.p010 = 0.0 | mass #self.staggered_grid.p110 = 0.0 | mass #self.staggered_grid.p000 = 0.0 | mass #self.staggered_grid.p100 = 0.0 | mass #self.staggered_grid.p011 = 0.0 | mass #self.staggered_grid.p111 = 0.0 | mass #self.staggered_grid.p001 = 0.0 | mass #self.staggered_grid.p101 = 0.0 | mass self.normal_grid = Grid.create(self.gridcode.grid.shape, corner1 - corner0 + delta) particles = Particles(self.normal_grid.size) particles.mass = 0.0 | mass particles.x = self.grid.x.flatten() particles.y = self.grid.y.flatten() particles.z = self.grid.z.flatten() particles.vx = 0.0 | speed particles.vy = 0.0 | speed particles.vz = 0.0 | speed particles.radius = 0.0 | length self.nbodycode.particles.add_particles(particles) self.from_model_to_nbody = particles.new_channel_to( self.nbodycode.particles) self.nbodycode.commit_particles() self.particles = particles
def setup_simple_grid(self): test_grid = Grid.create((4, 3, 2), [1.0, 1.0, 1.0] | units.m) test_grid.rho = numpy.linspace(1.0, 2.0, num=24).reshape( test_grid.shape) | units.kg / units.m**3 test_grid.rhovx = test_grid.rho * (3.0 | units.m / units.s) test_grid.rhovy = test_grid.rho * (4.0 | units.m / units.s) test_grid.rhovz = test_grid.rho * (0.0 | units.m / units.s) test_grid.energy = test_grid.rho * ((1.0 | (units.m / units.s)**2) + 0.5 * (5.0 | units.m / units.s)**2) return test_grid
def test_evolve_w_plot(sysfac, tend=1. | units.hour, dt=3600. | units.s, dtplot=None): sys = sysfac() if dtplot is None: dtplot = dt pyplot.ion() f = pyplot.figure(figsize=(10, 10)) pyplot.show() i = 0 Lx = sys.parameters.Lx grid = Grid.create((400, 400), (Lx, Lx)) dx, dy = grid.cellsize() Lx = Lx.value_in(1000 * units.km) x = grid.x.flatten() y = grid.y.flatten() while sys.model_time < tend - dtplot / 2: i = i + 1 sys.evolve_model(sys.model_time + dtplot, dt=dt) psi, dpsi = sys.get_psi_dpsidt(dx + 0. * x, x, y) psi = psi.reshape(grid.shape) # psi=sys.grid[:,:,0].psi f.clf() f1 = pyplot.subplot(111) f1.imshow(psi.transpose() / psi.max(), vmin=0, vmax=1, extent=[0, Lx, 0, Lx], origin="lower") f1.set_xlabel("x (x1000 km)") pyplot.draw() pyplot.savefig("test_bc.png") if i % 100 == 25: print("wait") raw_input() print("done") raw_input()
def new_grid(self): density = mass / length**3 density = density momentum = speed * density energy = mass / (time**2 * length) grid = Grid.create(self.dimensions_of_mesh, (10.0, 10.0, 10.0) | length) grid.rho = self.rho_medium grid.rhovx = 0.0 | momentum grid.rhovy = 0.0 | momentum grid.rhovz = 0.0 | momentum grid.energy = 0.0 | energy return grid
def new_grid(): grid = Grid.create(DIMENSIONS_OF_MESH, [1, 1, 1] | length) clear_grid(grid) return grid
def new_grid(self): grid = Grid.create(self.dimensions_of_mesh, [1,1,1] | length) self.clear_grid(grid) return grid
def new_grid(): grid = Grid.create(DIMENSIONS_OF_MESH, [1,1,1] | length) clear_grid(grid) return grid
def create_grid(*arg): grid = Grid.create(*arg) grid.add_vector_attribute("momentum", ["rhovx", "rhovy", "rhovz"]) return grid
def new_grid(self): grid = Grid.create(self.dimensions_of_mesh, [1, 1, 1] | length) self.clear_grid(grid) return grid
def initialize_grid(grid, magentic_field_grid): temp = [1, 1, 1] | length l = 1.0 | length A_grid = Grid.create(DIMENSIONS_OF_A_MESH, temp + grid.cellsize() * [2, 2, 0]) A_grid.position -= A_grid.cellsize() * [0.5, 0.5, 0.0] A_grid.Az = (B0 * l / (4.0 * PI) * numpy.cos(4.0 * PI / l * A_grid.x) + B0 * l / (2.0 * PI) * numpy.cos(2.0 * PI / l * A_grid.y)) assert grid.cellsize().x == A_grid.cellsize().x assert grid.cellsize().y == A_grid.cellsize().y # assert grid.dim[0] == A_grid.dim[0] magentic_field_grid.B1i = (A_grid.Az[:-1, 1:, ...] - A_grid.Az[:-1, :-1, ...]) / A_grid.cellsize().y magentic_field_grid.B2i = ( -(A_grid.Az[1:, :-1, ...] - A_grid.Az[:-1, :-1, ...]) / A_grid.cellsize().x) magentic_field_grid.B3i = 0.0 | magnetic_field # magentic_field_grid.B1i[-1,...,...] = magentic_field_grid.B1i[0,...,...] # magentic_field_grid.B2i[...,-1,...] = magentic_field_grid.B2i[...,0,...] magentic_field_grid.B1c = 0.0 | magnetic_field magentic_field_grid.B2c = 0.0 | magnetic_field magentic_field_grid.B3c = 0.0 | magnetic_field magentic_field_grid.divB = 0.0 | (magnetic_field / length) magentic_field_grid.B1c[:-1, ..., ...] = ( magentic_field_grid.B1i[:-1, ..., ...] + magentic_field_grid.B1i[1:, ..., ...]) * 0.5 magentic_field_grid.B2c[ ..., :-1, ...] = (magentic_field_grid.B2i[..., :-1, ...] + magentic_field_grid.B2i[..., 1:, ...]) * 0.5 magentic_field_grid.B3c = magentic_field_grid.B3i # magentic_field_grid.B1c[-1,...,...] = magentic_field_grid.B1c[0,...,...] # magentic_field_grid.B2c[...,-1,...] = magentic_field_grid.B2c[...,0,...] mu0 = 1.0 | (mass * length / time**2 / current**2) magentic_field_grid.divB[0:-1, 0:-1, ...] = ( (magentic_field_grid.B1i[:-1, :-1, ...] - magentic_field_grid.B1i[1:, :-1, ...]) / grid.cellsize().x + (magentic_field_grid.B2i[:-1, :-1, ...] - magentic_field_grid.B2i[:-1, 1:, ...]) / grid.cellsize().y) assert (abs(magentic_field_grid.divB[:-1, :-1, :-1]) < 1.0e-10 | magentic_field_grid.divB.unit).all() assert magentic_field_grid.B1i.unit == magnetic_field grid.rho = D0 grid.rhovx = -D0 * V0 * numpy.sin(2.0 * PI / l * grid.y) grid.rhovy = D0 * V0 * numpy.sin(2.0 * PI / l * grid.x) grid.rhovz = 0.0 | momentum print("sum rho vx:", grid.rhovx.sum()) print("sum rho vy:", grid.rhovy.sum()) print("sum rho vz:", grid.rhovz.sum()) grid.energy = (P0 / (GAMMA - 1) + 0.5 * (grid.rhovx**2 + grid.rhovy**2 + grid.rhovz**2) / grid.rho + 0.5 * (magentic_field_grid.B1c[:-1, :-1, :-1]**2 + magentic_field_grid.B2c[:-1, :-1, :-1]**2 + magentic_field_grid.B3c[:-1, :-1, :-1]**2) / mu0)
def hydro_plot(view, hydro_code, image_size, time, figname): """ Produce a series of images suitable for conversion into a movie. view: the (physical) region to plot [xmin, xmax, ymin, ymax] hydro_code: hydrodynamics code in which the gas to be plotted is defined image_size: size of the output image in pixels (x, y) time: current hydro code time """ if not HAS_MATPLOTLIB: return shape = (image_size[0], image_size[1], 1) size = image_size[0] * image_size[1] axis_lengths = [0.0, 0.0, 0.0] | units.m axis_lengths[0] = view[1] - view[0] axis_lengths[1] = view[3] - view[2] grid = Grid.create(shape, axis_lengths) grid.x += view[0] grid.y += view[2] speed = grid.z.reshape(size) * (0 | 1 / units.s) rho, rhovx, rhovy, rhovz, rhoe \ = hydro_code.get_hydro_state_at_point( grid.x.reshape(size), grid.y.reshape(size), grid.z.reshape(size), speed, speed, speed) min_v = 800.0 | units.km / units.s max_v = 3000.0 | units.km / units.s min_rho = 3.0e-9 | units.g / units.cm**3 max_rho = 1.0e-5 | units.g / units.cm**3 min_E = 1.0e11 | units.J / units.kg max_E = 1.0e13 | units.J / units.kg v_sqr = (rhovx**2 + rhovy**2 + rhovz**2) / rho**2 E = rhoe / rho log_v = numpy.log((v_sqr / min_v**2)) / numpy.log((max_v**2 / min_v**2)) log_rho = numpy.log((rho / min_rho)) / numpy.log((max_rho / min_rho)) log_E = numpy.log((E / min_E)) / numpy.log((max_E / min_E)) red = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_rho)).reshape(shape) green = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_v)).reshape(shape) blue = numpy.minimum(numpy.ones_like(rho.number), numpy.maximum(numpy.zeros_like(rho.number), log_E)).reshape(shape) alpha = numpy.minimum( numpy.ones_like(log_v), numpy.maximum(numpy.zeros_like(log_v), numpy.log( (rho / (10 * min_rho))))).reshape(shape) rgba = numpy.concatenate((red, green, blue, alpha), axis=2) pyplot.figure(figsize=(image_size[0] / 100.0, image_size[1] / 100.0), dpi=100) im = pyplot.figimage(rgba, origin='lower') pyplot.savefig(figname, transparent=True, dpi=100, facecolor='k', edgecolor='k') print("Saved hydroplot at time", time, "in file") print(' ', figname) pyplot.close()
def initialize_grid(grid, magentic_field_grid): temp = [1,1,1] | length l = 1.0 | length; A_grid = Grid.create(DIMENSIONS_OF_A_MESH, temp + grid.cellsize()*[2,2,0]) A_grid.position -= A_grid.cellsize()*[0.5,0.5,0.0] A_grid.Az = ( B0*l/(4.0*PI)*numpy.cos(4.0*PI/l*A_grid.x) + B0*l/(2.0*PI)*numpy.cos(2.0*PI/l*A_grid.y)) assert grid.cellsize().x == A_grid.cellsize().x assert grid.cellsize().y == A_grid.cellsize().y # assert grid.dim[0] == A_grid.dim[0] magentic_field_grid.B1i = (A_grid.Az[:-1,1:,...] - A_grid.Az[:-1,:-1,...])/A_grid.cellsize().y; magentic_field_grid.B2i = -(A_grid.Az[1:,:-1,...] - A_grid.Az[:-1,:-1,...])/A_grid.cellsize().x; magentic_field_grid.B3i = 0.0 | magnetic_field; #magentic_field_grid.B1i[-1,...,...] = magentic_field_grid.B1i[0,...,...] #magentic_field_grid.B2i[...,-1,...] = magentic_field_grid.B2i[...,0,...] magentic_field_grid.B1c = 0.0 | magnetic_field; magentic_field_grid.B2c = 0.0 | magnetic_field; magentic_field_grid.B3c = 0.0 | magnetic_field; magentic_field_grid.divB = 0.0 | (magnetic_field / length) magentic_field_grid.B1c[:-1,...,...] = (magentic_field_grid.B1i[:-1,...,...] + magentic_field_grid.B1i[1:,...,...])*0.5 magentic_field_grid.B2c[...,:-1,...] = (magentic_field_grid.B2i[...,:-1,...] + magentic_field_grid.B2i[...,1:,...])*0.5 magentic_field_grid.B3c = magentic_field_grid.B3i # magentic_field_grid.B1c[-1,...,...] = magentic_field_grid.B1c[0,...,...] # magentic_field_grid.B2c[...,-1,...] = magentic_field_grid.B2c[...,0,...] mu0 = 1.0 | (mass*length/time**2/current**2) magentic_field_grid.divB[0:-1,0:-1,...] = ( (magentic_field_grid.B1i[:-1,:-1,...] - magentic_field_grid.B1i[1:,:-1,...])/grid.cellsize().x + (magentic_field_grid.B2i[:-1,:-1,...] - magentic_field_grid.B2i[:-1,1:,...])/grid.cellsize().y ) assert (abs(magentic_field_grid.divB[:-1,:-1,:-1]) < 1.0e-10 | magentic_field_grid.divB.unit).all() assert magentic_field_grid.B1i.unit == magnetic_field grid.rho = D0 grid.rhovx = -D0*V0*numpy.sin(2.0*PI/l*grid.y) grid.rhovy = D0*V0*numpy.sin(2.0*PI/l*grid.x) grid.rhovz = 0.0 | momentum print("sum rho vx:", grid.rhovx.sum()) print("sum rho vy:",grid.rhovy.sum()) print("sum rho vz:",grid.rhovz.sum()) grid.energy = ( P0 / (GAMMA - 1) + 0.5 * (grid.rhovx**2 + grid.rhovy**2 + grid.rhovz**2)/grid.rho + 0.5 * (magentic_field_grid.B1c[:-1,:-1,:-1]**2 + magentic_field_grid.B2c[:-1,:-1,:-1]**2 + magentic_field_grid.B3c[:-1,:-1,:-1]**2)/mu0 )