fix the readme and doc examples

NVlabs · May 31, 2024 · 60ce89f · 60ce89f
1 parent 5277dd5
commit 60ce89f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -27,12 +27,13 @@ pip install --no-build-isolation earth2-grid
 
 ```
 >>> import earth2grid
+>>> import torch
 ... # level is the resolution
 ... level = 6
-... hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.PixelOrder.XY)
+... hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.XY())
 ... src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64)
-... z_torch = torch.as_tensor(z)
-... z_torch = torch.as_tensor(z)
+... z_torch = torch.cos(torch.deg2rad(torch.tensor(src.lat)))
+... z_torch = z_torch.broadcast_to(src.shape)
 >>> regrid = earth2grid.get_regridder(src, hpx)
 >>> z_hpx = regrid(z_torch)
 >>> z_hpx.shape

diff --git a/docs/usage.md b/docs/usage.md
@@ -3,16 +3,22 @@
 To use Earth2 Grid Utilities in a project
 
 ```
-import earth2grid
-# level is the resolution
-level = 6
-hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.PixelOrder.XY)
-src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64)
-z_torch = torch.as_tensor(z)
-z_torch = torch.as_tensor(z)
-regrid = earth2grid.get_regridder(src, hpx)
-z_hpx = regrid(z_torch)
-z_hpx.shape
-nside = 2**level
-reshaped = z_hpx.reshape(12, nside, nside)
+>>> import earth2grid
+>>> import torch
+... # level is the resolution
+... level = 6
+... hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.XY())
+... src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64)
+... z_torch = torch.cos(torch.deg2rad(torch.tensor(src.lat)))
+... z_torch = z_torch.broadcast_to(src.shape)
+>>> regrid = earth2grid.get_regridder(src, hpx)
+>>> z_hpx = regrid(z_torch)
+>>> z_hpx.shape
+torch.Size([49152])
+>>> nside = 2**level
+... reshaped = z_hpx.reshape(12, nside, nside)
+... lat_r = hpx.lat.reshape(12, nside, nside)
+... lon_r = hpx.lon.reshape(12, nside, nside)
+>>> reshaped.shape
+torch.Size([12, 64, 64])
 ```