do a safer divide when density is zero in consFlux (#297)

Author: zhichen3

pyro/compressible/riemann.py

@@ -1130,8 +1130,21 @@ def consFlux(idir, coord_type, gamma,
 
     F = np.zeros_like(U_state)
 
-    u = U_state[..., ixmom] / U_state[..., idens]
-    v = U_state[..., iymom] / U_state[..., idens]
+    if U_state.ndim == 1:
+        u = 0.0
+        v = 0.0
+        if U_state[idens] != 0.0:
+            u = U_state[ixmom] / U_state[idens]
+            v = U_state[iymom] / U_state[idens]
+
+    else:
+        u = np.zeros_like(U_state[..., ixmom])
+        v = np.zeros_like(U_state[..., iymom])
+        for i in range(U_state.shape[0]):
+            for j in range(U_state.shape[1]):
+                if U_state[i, j, idens] != 0.0:
+                    u[i, j] = U_state[i, j, ixmom] / U_state[i, j, idens]
+                    v[i, j] = U_state[i, j, iymom] / U_state[i, j, idens]
 
     p = (U_state[..., iener] - 0.5 * U_state[..., idens] * (u * u + v * v)) * (gamma - 1.0)
 

pyro/compressible/simulation.py

@@ -52,12 +52,18 @@ def cons_to_prim(U, gamma, ivars, myg):
     q = myg.scratch_array(nvar=ivars.nq)
 
     q[:, :, ivars.irho] = U[:, :, ivars.idens]
-    q[:, :, ivars.iu] = U[:, :, ivars.ixmom]/U[:, :, ivars.idens]
-    q[:, :, ivars.iv] = U[:, :, ivars.iymom]/U[:, :, ivars.idens]
-
-    e = (U[:, :, ivars.iener] -
-         0.5*q[:, :, ivars.irho]*(q[:, :, ivars.iu]**2 +
-                                  q[:, :, ivars.iv]**2))/q[:, :, ivars.irho]
+    q[:, :, ivars.iu] = np.divide(U[:, :, ivars.ixmom], U[:, :, ivars.idens],
+                                  out=np.zeros_like(U[:, :, ivars.ixmom]),
+                                  where=(U[:, :, ivars.idens] != 0.0))
+    q[:, :, ivars.iv] = np.divide(U[:, :, ivars.iymom], U[:, :, ivars.idens],
+                                  out=np.zeros_like(U[:, :, ivars.iymom]),
+                                  where=(U[:, :, ivars.idens] != 0.0))
+
+    e = np.divide(U[:, :, ivars.iener] - 0.5 * q[:, :, ivars.irho] *
+                  (q[:, :, ivars.iu]**2 + q[:, :, ivars.iv]**2),
+                  q[:, :, ivars.irho],
+                  out=np.zeros_like(U[:, :, ivars.iener]),
+                  where=(U[:, :, ivars.idens] != 0.0))
 
     q[:, :, ivars.ip] = eos.pres(gamma, q[:, :, ivars.irho], e)