switch to julia

filtering.jl

+module Filtering
+
+export imfilter, imfilter!, trim_kernel
+
+@inline function _shift_range(range, shift::Integer)
+  Base.UnitRange(first(range)+shift : last(range)+shift)
+end
+
+## TODO: allow border to be tuple
+#function _innershift_indices(a::Array{T, d}, shift::NTuple{d, Integer}, border = abs(maximum(shift))) where {T, d}
+#  abs(maximum(shift)) <= border || throw("cannot shift by $shift with border $border")
+#  range(i) = (1 + shift[i] + border):(lastindex(a, i) + shift[i] - border)
+#  ntuple(range, Val(ndims(a)))
+#end
+#
+#function innershift_view(a::Array{T, d}, shift::NTuple{d, Integer} = ntuple(i -> 0, d), border = abs(maximum(shift))) where {T, d}
+#  view(a, _innershift_indices(a, shift, border)...)
+#end
+
+@inline function _bidx_to_range(range::AbstractUnitRange, bidx::Integer)
+  if bidx > 0
+    return last(range):last(range)
+  elseif bidx < 0
+    return first(range):first(range)
+  else
+    return first(range)+1:last(range)-1
+  end
+end
+
+
+# we return Base.Slice since that is what `OffsetArray` uses for `axes()`
+function _trim_kernel_axes(kaxes, bidx)
+  ntuple(Val(length(kaxes))) do i
+    if bidx[i] < 0
+      return Base.Slice(0:last(kaxes[i]))
+    elseif bidx[i] > 0
+      return Base.Slice(first(kaxes[i]):0)
+    else
+      return Base.Slice(kaxes[i])
+    end
+  end
+end
+
+"""
+    trim_kernel(kernel, bidx)
+
+Trim `kernel` so that the resulting kernel is usable on the boundary
+indicated by `bidx`
+"""
+function trim_kernel(kernel, bidx)
+  dst_axes = _trim_kernel_axes(axes(kernel), bidx)
+  out = similar(kernel, dst_axes)
+  range = CartesianIndices(out)
+  copyto!(out, range, kernel, range)
+end
+
+
+
+# TODO: weirdly Base.Slice instead of Base.UnitRange leads to axes that are not
+# compatible when broadcasting, figure out why!
+function imfilter!(dst, img, kern, slice=_bidx_to_range.(axes(img), ntuple(i->0, ndims(img))))
+  ndims(dst) == ndims(img) == ndims(kern) == length(slice) ||
+      throw(ArgumentError("dst, image, kernel or slice dimensions disagree ($(ndims(dst)) vs $(ndims(img)) vs $(ndims(kern)) vs $(length(slice)))"))
+  axes(dst) == axes(img) || throw(ArgumentError("output and input image axes disagree"))
+  all(size(kern) .<= size(img)) || throw(ArgumentError("kernel bigger than image"))
+  last.(slice) .+ last.(axes(kern)) <= last.(axes(img)) || first.(slice) .+ first.(axes(kern)) >= first.(axes(img)) ||
+      throw(ArgumentError("kern and/or slice out of bounds for image"))
+
+  dstslice = view(dst, UnitRange.(slice)...)
+  for i in CartesianIndices(kern)
+    iszero(kern[i]) && continue
+    dstslice .+= kern[i] .* view(img, (_shift_range.(slice, Tuple(i)))...)
+  end
+  dst
+end
+
+"""
+    filter(kernelf, img)
+
+Filter `img` using a kernel function `kernelf(bidx) = kernel` specifying a
+kernel for each type of node.
+"""
+function imfilter(kernelf::Function, img)
+  out = zeros(axes(img))
+  for bidx in Iterators.product(ntuple(i -> -1:1, ndims(img))...)
+    range = _bidx_to_range.(axes(img), bidx)
+    imfilter!(out, img, kernelf(bidx), range)
+  end
+  out
+end
+
+
+"""
+    imfilter(img, kern)
+
+Filter `img` by `kern` using correlation.
+Boundary is zeroed out.
+"""
+imfilter(img, kern) = imfilter(bidx -> trim_kernel(kern, bidx), img)
+# TODO: use keyword arguments
+imfilter_inner(img, kern) = imfilter(bidx -> all(bidx .== 0) ? kern : zeros(ntuple(i -> 0, ndims(img))), img)
+
+end

test.jl

+module FDM
+
+using StaticArrays
+using SparseArrays
+using OffsetArrays
+using LinearMaps
+using IterativeSolvers
+
+include("filtering.jl")
+
+"Finite Differences"
+module FD
+
+using OffsetArrays
+
+const forward = OffsetVector([-1, 0, 1], -2)
+const central = OffsetVector([-1, 0, 1], -2)
+function dirfd(fd, ndims, dir)
+  dir <= ndims || throw(ArgumentError("need dir <= ndims!"))
+  dfd = OffsetArray(zeros(ntuple(i -> i == dir ? length(fd) : 1, ndims)), ntuple(i -> i == dir ? fd.offsets[1] : -1, ndims))
+  # Workaround: LinearIndices(::OffsetVector) are not 1-based, so copyto! chockes
+  copyto!(dfd, fd.parent)
+end
+
+end
+
+using .Filtering
+using .FD
+
+
+
+
+# aka `StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
+const FDStepRange = typeof(0:0.1:1)
+
+struct Grid{d}
+  domain::NTuple{d, FDStepRange}
+end
+
+Base.getindex(grid::Grid, idx...) = SVector(getindex.(grid.domain, idx))
+
+
+"""
+    neuman_kernel(kernel, bidx)
+
+Modify `kernel` to be usable on a boundary `bixd` in a neumannn boundary setting.
+
+Ghost nodes are eliminated through central finite differences. Currently only
+the basic 5-star Laplace is supported.
+"""
+function neumann_kernel(kernel, bidx)
+  out = copy(kernel)
+  for (i, bi) in enumerate(bidx)
+    colrange = (i == j ? Colon() : 0 for j in 1:ndims(kernel))
+    out[colrange...] .+= bi * FD.forward
+  end
+  out /= 2 ^ count(!iszero, bidx)
+  trim_kernel(out, bidx)
+end
+
+function divergence(field::AbstractArray{T, d}) where {T, d}
+  size(field, d) == d - 1 || throw(ArgumentError("need a d-vectorfield on a d-grid"))
+  dst = zeros(size(field))
+  for k in 1:d-1
+    kern = FD.dirfd(FD.central, d-1, k)
+    imfilter!(selectdim(dst, d, k), selectdim(field, d, k), kern)
+  end
+  dropdims(sum(dst, dims=d), dims=d)
+end
+
+function gradient(img::AbstractArray)
+  d = ndims(img)
+  dst = zeros(size(img)..., d)
+  for k in 1:d
+    kern = FD.dirfd(FD.central, d, k)
+    imfilter!(selectdim(dst, d+1, k), img, kern)
+  end
+  dst
+end
+
+
+function apply(kern, flatimg, size)
+  img = reshape(flatimg, size)
+  out = imfilter(kern, img)
+  reshape(out, prod(size))
+end
+
+function solve(pde, shape::NTuple{d, Integer}) where d
+  ndims(pde.stencil) == d || throw("dimension mismatch")
+  length(pde.f) == prod(shape) || throw("bounds mismatch")
+
+  A = LinearMap{Float64}(u -> apply(pde.stencil, u, shape), prod(shape), issymmetric=true, isposdef=true)
+
+  u = zeros(prod(shape))
+  cg!(u, A, pde.f)
+  return u
+end
+
+end
+
+using StaticArrays
+using OffsetArrays
+using LinearAlgebra
+
+st = MMatrix{3, 3}([0 -1 0; -1 4 -1; 0 -1 0])
+st = OffsetArray(st, (-2, -2))
+
+img = rand(1000, 1000)
+
+n = 100
+d = 2
+
+β = 1e-3
+A = I
+τ = 0.5
+α = 1
+
+f = rand(n, n)
+p = rand(n, n, d)
+
+x = (A'*A + β*I) \ - FDM.divergence(p) - A'*f
+w = FDM.gradient(x)
+
+p = (p - τ*w) ./ (1 .+ τ/α * abs.(w))
+
+
+#@benchmark solve((stencil=st, f=zeros(10000)), (100, 100))
+
+#u = zeros(100, 100)
+
+#A = LinearMap{Float64}(u -> apply(st, u, size(uinner)), length(uinner), issymmetric=true, isposdef=true)
+
+#cg!(uflat, A, zeros(length(uinner)))
+
+
+#A = sparse(A)
+
+#@benchmark cg($A, zeros(10000))
+
+

test.py

+import numpy as np
+from scipy.sparse import diags, coo_matrix
+from scipy.sparse.linalg import spsolve, cg
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.pyplot as plt
+
+# TODO: use np.indices, np.mgrid, np.meshgrid, etc.
+
+
+domain = np.array([100, 100, 100])
+d = domain.size
+N = np.prod(domain)
+
+
+# laplace stencil in d dimensions
+def laplace(d):
+    st1 = np.array([1, -2, 1])
+    s = st1.size
+
+    st = np.zeros(np.full(d, s))
+
+    idx = np.empty(d, dtype=slice)
+    idx.fill(st1.size // 2)
+    idx[0] = slice(None)
+
+    for k in range(d):
+        st[tuple(idx)] += st1
+        idx = np.roll(idx, 1)
+
+    return st
+
+stencil = laplace(d)
+
+print(stencil)
+
+
+#stencil = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
+
+
+u = np.zeros(N)
+f = 1 * np.ones(N)
+
+# boundary nodes (first and last planes are special)
+bnd = np.full(N, False)
+for nd in np.cumprod(np.flip(domain[1:], axis=0)):
+    print(f"nd = {nd}")
+    bnd = bnd | (np.arange(0, N) % nd == 0)
+    bnd = bnd | (np.arange(0, N) % -nd == -1)
+bnd = bnd | (np.arange(0, N) // domain[-1] == 0)
+bnd = bnd | (np.arange(0, N) // domain[-1] == N // domain[-1] - 1)
+
+bnd_cnt = np.count_nonzero(bnd)
+free_cnt = np.prod(domain) - bnd_cnt
+
+
+# TODO: set u boundary data
+#u[bnd] = np.arange(0, N)[bnd]
+
+print(f"boundary nodes: {np.count_nonzero(bnd)}")
+print(f"free nodes: {np.count_nonzero(~bnd)}")
+
+elstrides = np.flip(np.cumprod(np.flip(np.append(domain[1:], [1]), 0)), 0)
+diagonals = []
+offsets = []
+
+
+center = (np.array(stencil.shape) - 1) // 2
+for index, value in np.ndenumerate(stencil):
+    if value == 0:
+        continue
+    offset = np.array(index) - center
+    offsets.append(np.dot(elstrides, offset))
+    diagonals.append(value)
+
+print(diagonals, offsets)
+A = diags(diagonals, offsets, np.full(2, np.prod(domain))).tocoo()
+
+f = f - A @ u
+
+def filter_free(bnd, A):
+    assert A.shape[0] == A.shape[1]
+    print("filtering ...", end=" ", flush=True)
+
+    # filter out boundary nodes
+    spfree = ~bnd[A.row] & ~bnd[A.col]
+    rows, cols, data = A.row[spfree], A.col[spfree], A.data[spfree]
+
+    # correct indices
+    offset = -np.cumsum(bnd)
+    rows = rows + offset[rows]
+    cols = cols + offset[cols]
+
+    nf = A.shape[0] - np.count_nonzero(bnd)
+    Af = coo_matrix((data, (rows, cols)), shape=(nf, nf))
+    print("finished!")
+    return Af
+
+Af = filter_free(bnd, A).tocsr()
+
+#Af = A[np.ix_(~bnd, ~bnd)]
+ff = f[~bnd]
+
+print("solving ...", end=" ", flush=True)
+u[~bnd] = spsolve(Af, ff)
+#u[~bnd] = cg(Af, ff)
+print("finished!")
+
+
+x, y = np.meshgrid(np.arange(0, domain[1]), np.arange(0, domain[0]))
+
+print(x.shape, y.shape)
+z = u.reshape(domain)
+
+fig = plt.figure()
+ax = fig.add_subplot(111, projection='3d')
+ax.plot_surface(x, y, z)
+plt.show()
+#
+#
+#print(u.reshape(domain))
+#
+#print(np.diagflat(np.full((n), value)))
+
+
+#np.roll(p, -1) + 2*p + np.roll(
+
+
+
+
+#s = np.sin(2 * np.pi * t)
+#
+#plt.plot(t, s)
+#plt.show()
+