help_text = '''

'''

basename = os.path.splitext(filename)[0]

bvecpath = f'{basename}.bvec'

bvalpath = f'{basename}.bval'

if os.path.exists(bvecpath):

bvec = np.loadtxt(bvecpath).T

if os.path.exists(bvalpath):

bval = np.loadtxt(bvalpath)[:,None]

'''

basename = os.path.splitext(filename)[0]

bvecpath = f'{basename}.bvec'

bvalpath = f'{basename}.bval'

np.savetxt(bvecpath, bvec.T, fmt='%.6f')

'''

gradv = []

gradb = []

with open(filename, 'r') as f:

for line in f:

line = line.strip().split()

v = list(map(float, line[:3]))

b = float(line[3])

gradv.append(v)

gradb.append(b)

gradv = np.array(gradv)

gradb = np.array(gradb)[:,None]

'''

'''

with open(filename, 'w') as f:

for v, b in zip(gradv, gradb):

line = f'{v[0]:.7g} {v[1]:.7g} {v[2]:.7g} {b.item():g}\n'

'''

metadata = {}

vectors = []

vpattern = re.compile(r'Vector\[\d+\]\s*=\s*\(([-\d\s.,]+)\)')

mpattern = re.compile(r'^\s*([a-zA-Z]+)\s*=\s*([^\s#]+)')

with open(filename, 'r') as f:

for line in f:

line = line.strip()

if not line or line.startswith('#'):

continue

mmatch = mpattern.match(line)

if mmatch:

key = mmatch.group(1).strip()

value = mmatch.group(2).strip()

metadata[key] = value

continue

vmatch = vpattern.search(line)

if vmatch:

vals = vmatch.group(1).replace(',', ' ').split()

vec = [float(n) for n in vals]

vectors.append(vec)

vectors = np.array(vectors)

'''

if metadata is None:

metadata = {'CoordinateSystem': 'xyz', 'Normalization': 'none'}

with open(filename, 'w') as f:

f.write('[directions={}]\n'.format(len(vectors)))

for key, value in metadata.items():

f.write(f'{key} = {value}\n')

for i, v in enumerate(vectors):

'''

with open(filename) as f:

lines = [line.strip() for line in f if not line.strip().isdigit()]

vectors = np.array([list(map(float, line.split())) for line in lines])

'''

with open(filename, 'w') as f:

f.write(f"{len(vectors)}\n")

for v in vectors:

'''

vectors = []

bvalues = []

with open(filename, 'r') as f:

comment = f.readline()

line = comment.strip().split()

if len(line) == 4 and \

all(part.replace('.', '', 1).replace('-', '', 1).isdigit() for part in line):

vector = list(map(float, line[:3]))

bvalue = int(line[3]) # keep separate if we need to change to int in the future

vectors.append(vector)

bvalues.append(bvalue)

comment = None

for line in f:

line = line.strip().split()

vector = list(map(float, line[:3]))

bvalue = float(line[3]) # keep separate if we need to change to int in the future

vectors.append(vector)

bvalues.append(bvalue)

vectors = np.array(vectors)

bvalues = np.array(bvalues)[:,None]

'''

if bvalues[0] != 0:

raise ValueError('For Philips, the first b-value must be 0.')

if np.unique(vectors, axis=0).shape[0] != vectors.shape[0]:

raise ValueError('For Philips, duplicate vectors are not allowed, not even for b=0.')

with open(filename, 'w') as f:

if comment:

f.write(comment + '\n')

for vector, bvalue in zip(vectors, bvalues):

'''

shells = []

vectors = []

with open(filename, 'r') as f:

for line in f:

if line.strip().startswith('#') or not line.strip():

continue

parts = line.strip().split()

shells.append(int(parts[0])) # shell number

vectors.append([float(parts[1]), float(parts[2]), float(parts[3])])

vectors = np.array(vectors)

shells = np.array(shells)[:,None]

shells = shells - shells.min() # we want the shells in python to start from 0

'''

shells = shells - shells.min() + 1 # indices in the file start from 1

with open(filename, 'w') as f:

f.write('#shell\tu_x\tu_y\tu_z\n')

for shell, vector in zip(shells, vectors):

'''

Ks = np.array(Ks)

K = np.sum(Ks)

W, shells = calc_weights(Ks, alpha) # precompute weights

vectors = random_vectors(K)

svectors = vectors.ravel(order='C') # they need be stacked for optimization

result = scopt.minimize(

cost_all_vectors,

svectors,

args = (K,W,),

method = 'SLSQP',

jac = grad_all_vectors,

constraints = [{'type' : 'eq',

'fun' : eq_constraints,

'args' : (K,)}],

options = {'maxiter': maxiter,

'ftol' : epsilon,

'disp' : False})

svectors = result.x

vectors = svectors.reshape((K,3))

vectors = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)

"""

# This function could be simpler, without iterating twice over range(S)

# but it's kept as this for clarity and symmetry with the paper.

K = np.sum(Ks) # this can be inferred from

S = len(Ks)

shells = np.concatenate([np.full(Ks[s],s) for s in range(S)]) # starts at 0

W = np.zeros((K, K))

# V1: repulsion within shell, Eqn 2

for s in range(S):

idxs = (shells == s)

intraW = alpha / (S * Ks[s]**2)

idxss = np.outer(idxs,idxs) - np.diag(idxs.astype(float))

W += idxss * intraW

# V2: repulsion vetween shells, Eqn 4

interW = (1 - alpha) / K**2

for s in range(S):

idxs = (shells == s)

for t in range(S):

if s != t:

idxt = (shells == t)

idxst = np.outer(idxs,idxt)

idxts = np.outer(idxt,idxs)

W += (idxst + idxts) * interW

shells = shells[:,None] # shells start counting from 0

'''

vectors = svectors.reshape((K,3))

vectors = vectors / np.linalg.norm(vectors, axis=1)[:, None]

# Equation 3 of the paper. For unit vectors, note that:

# ||u - v||^2 = 2 - 2.u.v

# ||u + v||^2 = 2 + 2.u.v

D = np.dot(vectors, vectors.T)

energy = (1 / (2 + 2*D + epsilon)) + (1 / (2 - 2*D + epsilon))

Energy = np.sum(W * energy)

'''

vectors = svectors.reshape((K,3))

vectors = vectors / np.linalg.norm(vectors, axis=1)[:, None]

num1 = vectors[:,None,:] - np.transpose(vectors[:,None,:], axes=(1,0,2))

num2 = vectors[:,None,:] + np.transpose(vectors[:,None,:], axes=(1,0,2))

D = np.dot(vectors, vectors.T)

den1 = (2 - 2*D)**2 + epsilon # ||u - u||^4

den2 = (2 + 2*D)**2 + epsilon # ||u + u||^4

pderiv = num1/den1[:,:,None] + num2/den2[:,:,None] # Eqn. 7 of the paper, partial derivatives for the 3 axes

grad = -2*np.sum(W[:,:,None]*pderiv, axis=1) # sum over axis 0 or 1 but not both (either u-v or v-u, but not both)

grad = grad.ravel(order='C')

'''

vectors = svectors.reshape((K,3), order='C')

eq = np.linalg.norm(vectors, axis=1) - 1.0

'''

_, Ks = np.unique(shells, return_counts=True)

S = len(Ks)

sphvec = cart2sph(vectors)[:,1:] # theta and phi only

idx = [] # assigned directions

seq = list(range(len(vectors))) # unassigned directions

while len(seq):

Ksc = np.array([np.sum(shells[idx] == s) for s in range(S)])

shell = next_shell(Ksc, Ks)

energy = np.full((len(seq),1), np.inf)

for i, sq in enumerate(seq):

if shells[sq] == shell:

energy[i] = cost_one_vector(sphvec[sq], shells[sq], sphvec[idx], shells[idx], np.ones((S,S)))

sidx = np.argmin(energy)

idx.append(seq[sidx])

seq.remove(seq[sidx])

vectors = vectors[idx]

shells = shells[idx]

'''

Ks = np.array(Ks)

K = np.sum(Ks)

S = len(Ks)

Ksc = np.zeros(S, dtype=int) # current directions per shell, will eventually match Ks

vectors = [] # spherical coordinates, only theta and phi

shells = []

W = np.full((S,S), 1-alpha)

np.fill_diagonal(W, alpha)

for k in range(K):

shell = next_shell(Ksc, Ks)

vector = find_one_vector(shell, vectors, shells, W, maxiter=maxiter)

vectors.append(vector)

shells.append(shell)

Ksc[shell] += 1

vectors = np.array(vectors) # spherical coordinates, theta and phi only

r = np.ones((len(vectors),1))

vectors = np.hstack((r, vectors)) # now with the radii

vectors = sph2cart(vectors) # cartesian coordinates

shells = np.array(shells)[:,None] # shells start counting from 0

'''

init_guess = cart2sph(random_vectors(1))[0,1:] # theta and phi only, as r = 1

result = scopt.minimize(

cost_one_vector,

init_guess,

args = (shell, vectors, shells, W),

method = 'L-BFGS-B',

bounds = [(0,np.pi), (0,2*np.pi)],

options = {'maxiter': maxiter,

'ftol' : epsilon,

'disp' : False})

vector = result.x # theta and phi only

'''

u = sph2cart(np.array((1, *vector)))

Energy = 0.0

for i, v in enumerate(vectors):

v = sph2cart(np.array((1, *v)))

diff = np.sum((u - v) ** 2) + epsilon

ssq = np.sum((u + v) ** 2) + epsilon

energy = 1/diff + 1/ssq # Eqn 3 of the paper

w = W[shell, shells[i]]

Energy += w * energy

'''

K = np.sum(Ks) # desired total number of directions

Kc = np.sum(Ksc) # current total number of directions

Ps = Ks / K # desired proportion of directions across shells

if Kc == 0:

deficits = Ps

else:

Pc = Ksc / Kc # current proportion

deficits = Ps - Pc

nextshell = np.argmax(deficits) # starts at 0

'''

if distrib in [0, 'uniform', 'uniformly', 'evenly', 'constant', 'constantly']:

Ps = np.ones(S)

elif distrib in [1, 'lin', 'linear', 'linearly']:

Ps = np.arange(1,S+1, dtype=float)

elif distrib in [2, 'quad', 'quadratic', 'quadratically']:

Ps = np.arange(1,S+1, dtype=float) ** 2

Ks0 = Ps/np.sum(Ps)*K

Ks = Ks0.astype(int)

while np.sum(Ks) < K:

Kdiff = Ks0 - Ks

Ks[np.argmax(Kdiff)] += 1

while np.sum(Ks) > K:

Kdiff = Ks0 - Ks

Ks[np.argmin(Kdiff)] -= 1

'''

vectors = np.random.randn(K,3)

vectors = vectors / np.linalg.norm(vectors, axis=1)[:,None]

'''

r = np.linalg.norm(cart, axis=1)

theta = np.arccos(cart[:,2] / r)

theta = np.where(r == 0, 0, theta) # rare case in which r = 0

phi = np.mod(np.arctan2(cart[:,1], cart[:,0]), 2*np.pi)

sph = np.column_stack((r, theta, phi))

'''

r, theta, phi = sph.T

cart = np.column_stack((

r * np.sin(theta) * np.cos(phi),

r * np.sin(theta) * np.sin(phi),

r * np.cos(theta) ))

'''

rA, cA = A.shape

rB, cB = B.shape

rM = rA + rB

if cA != cB:

raise ValueError('The two matrices must have the number of columns.')

M = np.empty((rM,cA))

if rA == rB:

M[0::2] = B

M[1::2] = A

return M

elif rA > rB:

nI = int(np.ceil(rM/rB)*rB)

idx = np.reshape(np.arange(nI), newshape=(rB,int(nI/rB)))

idxA = np.ravel(idx[:,1:], order='C')[:rA]

idxB = idx[:,0]

elif rA < rB:

nI = int(np.ceil(rM/rA)*rA)

idx = np.reshape(np.arange(nI), newshape=(rA,int(nI/rA)))

idxA = idx[:,0]

idxB = np.ravel(idx[:,1:], order='C')[:rB]

M[idxA,:] = A

M[idxB,:] = B

'''

if K0 == 0:

return vectors, bvalues

if unique:

vectors0 = random_vectors(K0)

else:

vectors0 = np.zeros((K0, 3))

bvalues0 = np.zeros((K0, 1))

if where == 'start':

vectors = np.concatenate((vectors0, vectors), axis=0)

bvalues = np.concatenate((bvalues0, bvalues), axis=0)

elif where == 'end':

vectors = np.concatenate((vectors, vectors0), axis=0)

bvalues = np.concatenate((bvalues, bvalues0), axis=0)

elif where == 'interspersed':

vectors = intersperse_rows(vectors, vectors0)

bvalues = intersperse_rows(bvalues, bvalues0)

'''

idx = bvalues != 0

bvalues = bvalues[idx]

vectors = vectors[idx]

shells = shells[idx]

'''

bmax = bvalues.max()

if bmax == 0:

vectors = np.full(vectors.shape,0)

else:

vectors = np.sqrt(bvalues/bmax)*vectors

'''

norm = np.linalg.norm(vectors, axis=1)

bvalues = np.round(norm**2 * bmax).astype(int)[:,None]

idx = norm != 0

vectors[idx,:] /= norm[idx,None]

'''

def subdivide(vtx, fac):

vtxn = vtx.tolist()

facn = []

edg2mid = {}

next_index = len(vtx)

def get_midpoint(i, j):

nonlocal next_index

if i > j:

i, j = j, i

key = (i,j)

if key not in edg2mid:

v_i, v_j = vtx[i], vtx[j]

mid = (v_i + v_j) / 2

mid /= np.linalg.norm(mid)

vtxn.append(mid)

edg2mid[key] = next_index

next_index += 1

return edg2mid[key]

for f in fac:

a, b, c = f

m_ab = get_midpoint(a, b)

m_bc = get_midpoint(b, c)

m_ca = get_midpoint(c, a)

facn.extend([

[a, m_ab, m_ca],

[m_ab, b, m_bc],

[m_ca, m_bc, c],

[m_ab, m_bc, m_ca]])

return np.array(vtxn), np.array(facn)

g = (1 + np.sqrt(5)) / 2 # Golden ratio

vtx = np.array([

[0, 1, g], [0, -1, g], [0, 1, -g], [0, -1, -g],

[1, g, 0], [-1, g, 0], [1, -g, 0], [-1, -g, 0],

[g, 0, 1], [g, 0, -1], [-g, 0, 1], [-g, 0, -1]

])

norm = np.linalg.norm(vtx, axis=1, keepdims=True)

vtx /= norm

fac = np.array([

[ 0, 1, 8], [ 0, 10, 1], [ 2, 9, 3], [ 2, 3, 11],

[ 1, 7, 6], [ 3, 6, 7], [ 5, 0, 4], [ 2, 5, 4],

[ 6, 9, 8], [ 8, 9, 4], [ 7, 10, 11], [ 5, 11, 10],

[ 1, 6, 8], [ 0, 8, 4], [10, 7, 1], [10, 0, 5],

[ 6, 3, 9], [ 7, 11, 3], [ 2, 4, 9], [ 2, 11, 5] ])

for _ in range(n):

vtx, fac = subdivide(vtx, fac)

'''

vtx = []

vtx.append([0.0, 0.0, 1.0]) # North pole

for i in range(1, nlat): # Rings between poles

theta = i * np.pi / nlat

for j in range(nlon):

phi = j * 2 * np.pi / nlon

x = np.sin(theta) * np.cos(phi)

y = np.sin(theta) * np.sin(phi)

z = np.cos(theta)

vtx.append([x, y, z])

vtx.append([0.0, 0.0, -1.0]) # South pole

# Faces (all triangular)

fac = []

for j in range(nlon): # North pole

v0 = 0 # North pole

v1 = 1 + j

v2 = 1 + (j + 1) % nlon

fac.append([v0, v1, v2])

for k in range(1, nlat - 1): # Middle section

for j in range(nlon):

v0 = 1 + (k - 1) * nlon + j

v1 = 1 + (k - 1) * nlon + (j + 1) % nlon

v2 = 1 + k * nlon + (j + 1) % nlon

v3 = 1 + k * nlon + j

# Split quad into two triangles

fac.append([v0, v1, v2])

fac.append([v0, v2, v3])

N = 1 + (nlat - 1) * nlon # South pole index

for j in range(nlon): # South pole

v0 = N

v1 = 1 + (nlat - 2) * nlon + (j + 1) % nlon

v2 = 1 + (nlat - 2) * nlon + j

fac.append([v0, v1, v2])

vtx = np.array(vtx)

fac = np.array(fac)

'''

# Vertex indices of interest

zpos = vtx[:,2] >= 0

newvtx = vtx[zpos]

# Map from old to new indices

zidx = np.where(zpos)[0]

newidx = np.full(len(vtx), -1, dtype=int)

newidx[zidx] = np.arange(len(zidx))

# Select faces where all vertices have z >= 0

fpos = np.all(zpos[fac], axis=1)

fidx = fac[fpos]

newfac = newidx[fidx]

filename=None, colorby='shell', sphere='uv',

style='quiver', reproject=False): # =======================

"""

import matplotlib

import matplotlib.pyplot as plt

import matplotlib.cm as cm

import matplotlib.colors as mcolors

from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# Do not open a window if saving to a file

if filename is not None:

matplotlib.use('agg') # Force Agg backend, override existing

# Here we want the shell indices to start at 1

if shells is None and bvalues is not None:

if not reproject:

vectors = scale_vectors(vectors, bvalues)

_, shells = np.unique(bvalues, return_inverse=True)

shells = shells + 1

bmax = bvalues.max()

bvalues1 = bvalues / bmax # scaled to between 0 and 1

elif shells is not None and bvalues is None:

shells = shells - shells.min() + 1

bmax = None

bvalues1 = shells / shells.max() # fake b-values, and scaled to between 0 and 1

if not reproject:

vectors = vectors * shells/shells.max()

elif shells is None and bvalues is None:

shells = np.ones(vectors.shape[0])

bmax = None

bvalues1 = np.ones(vectors.shape[0])

else:

raise ValueError('Must provide either "shells" or "bvalues", not both')

uB1 = np.unique(bvalues1)

uS = np.unique(shells)

S = len(uS)

K = len(vectors)

# We only need one hemisphere of the shells; let's flip so that they

# are all on the +z hemisphere, i.e., above the "equator".

idx = vectors[:,2] < 0

vectors[idx] = -vectors[idx]

# Custom colormap . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

if 'multishell' not in plt.colormaps():

colors = [(.95, .05, .05),

(.05, .85, .05),

(.05, .05, .95)]

cmap = mcolors.LinearSegmentedColormap.from_list('multishell', colors, N=2**10)

matplotlib.colormaps.register(cmap)

cmap = plt.get_cmap('multishell')

if colorby == 'shell':

C = S

Cidx = shells

uCidx = uS

elif colorby == 'acquisition':

C = K

Cidx = np.arange(K) + 1

uCidx = Cidx

colors = [cmap(c / (C - 1)) if C > 1 else cmap(0.5) for c in range(C)]

cmap = mcolors.ListedColormap(colors)

bounds = np.concatenate([[uCidx[0] - 0.5],

(uCidx[:-1] + uCidx[1:]) / 2,

[uCidx[-1] + 0.5]])

norm = mcolors.BoundaryNorm(bounds, cmap.N)

# Set up axes . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

fig = plt.figure(figsize=(10, 8))

ax = fig.add_subplot(111, projection='3d')

if style == 'qspace':

ax.grid(True)

ax.xaxis.pane.set_alpha(.5)

ax.yaxis.pane.set_alpha(.5)

ax.zaxis.pane.set_alpha(.5)

ax.xaxis.line.set_visible(True)

ax.yaxis.line.set_visible(True)

ax.zaxis.line.set_visible(True)

else:

ax.grid(False)

ax.xaxis.pane.set_alpha(0)

ax.yaxis.pane.set_alpha(0)

ax.zaxis.pane.set_alpha(0)

ax.xaxis.line.set_visible(False)

ax.yaxis.line.set_visible(False)

ax.zaxis.line.set_visible(False)

ax.set_xticks([])

ax.set_yticks([])

ax.set_zticks([])

ax.set_xlim(-1,1)

ax.set_ylim(-1,1)

ax.set_zlim(0,1)

ax.set_box_aspect([1,1,.5])

ax.set_proj_type('ortho')

# Directions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

if style == 'points':

sc = ax.scatter(vectors[:,0], vectors[:,1], vectors[:,2],

c=Cidx.ravel(order='C'), cmap=cmap, norm=norm, alpha=1, s=5)

elif style == 'qspace':

sphere = None

sc = ax.scatter(vectors[:,0], vectors[:,1], vectors[:,2],

c=Cidx.ravel(order='C'), cmap=cmap, norm=norm, alpha=1, s=5)

elif style == 'quiver':

for i, vec in enumerate(vectors):

if reproject:

arrow_length_ratio = .1

else:

if bvalues is None:

arrow_length_ratio = .1/bvalues1[i]

else:

arrow_length_ratio = .1/np.sqrt(bvalues1[i])

ax.quiver(0, 0, 0, vec[0], vec[1], vec[2],

color=cmap(norm(Cidx[i])), alpha=1, linewidth=1.5,

arrow_length_ratio=arrow_length_ratio)

elif style == 'rings':

for i, vec in enumerate(vectors):

if reproject:

r = 0.05

else:

if bvalues is None:

r = 0.05 * bvalues1[i]

else: