63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
def make_bounds(dirs, rivers):
|
|
(Y, X) = dirs.shape
|
|
bounds_h = np.zeros((Y, X-1), dtype='i4')
|
|
bounds_v = np.zeros((Y-1, X), dtype='i4')
|
|
|
|
bounds_v += (rivers * (dirs==1))[:-1,:]
|
|
bounds_h += (rivers * (dirs==2))[:,:-1]
|
|
bounds_v -= (rivers * (dirs==3))[1:,:]
|
|
bounds_h -= (rivers * (dirs==4))[:,1:]
|
|
|
|
return bounds_h, bounds_v
|
|
|
|
def get_fixed(dirs):
|
|
borders = np.zeros(dirs.shape, dtype='?')
|
|
borders[-1,:] |= dirs[-1,:]==1
|
|
borders[:,-1] |= dirs[:,-1]==2
|
|
borders[0,:] |= dirs[0,:]==3
|
|
borders[:,0] |= dirs[:,0]==4
|
|
|
|
donors = np.zeros(dirs.shape, dtype='?')
|
|
donors[1:,:] |= dirs[:-1,:]==1
|
|
donors[:,1:] |= dirs[:,:-1]==2
|
|
donors[:-1,:] |= dirs[1:,:]==3
|
|
donors[:,:-1] |= dirs[:,1:]==4
|
|
return borders | ~donors
|
|
|
|
def twist(bounds_x, bounds_y, fixed, d=0.1, n=5):
|
|
moveable = ~fixed
|
|
|
|
(Y, X) = fixed.shape
|
|
offset_x = np.zeros((Y, X))
|
|
offset_y = np.zeros((Y, X))
|
|
|
|
for i in range(n):
|
|
force_long = np.abs(bounds_x) * (1+np.diff(offset_x, axis=1))
|
|
force_trans = np.abs(bounds_y) * np.diff(offset_x, axis=0)
|
|
|
|
force_x = np.zeros((Y, X))
|
|
force_x[:,:-1] = force_long
|
|
force_x[:,1:] -= force_long
|
|
force_x[:-1,:]+= force_trans
|
|
force_x[1:,:] -= force_trans
|
|
|
|
force_long = np.abs(bounds_y) * (1+np.diff(offset_y, axis=0))
|
|
force_trans = np.abs(bounds_x) * np.diff(offset_y, axis=1)
|
|
|
|
force_y = np.zeros((Y, X))
|
|
force_y[:-1,:] = force_long
|
|
force_y[1:,:] -= force_long
|
|
force_y[:,:-1]+= force_trans
|
|
force_y[:,1:] -= force_trans
|
|
|
|
length = np.hypot(force_x, force_y)
|
|
length[length==0] = 1
|
|
coeff = d / length * moveable
|
|
offset_x += force_x * coeff
|
|
offset_y += force_y * coeff
|
|
|
|
return offset_x, offset_y
|