fzimmermann89/projection.py

## projection.py
import torch
from scipy.spatial.transform import Rotation
from typing import Callable, Literal
import einops
import itertools
from torch import Tensor


class MatrixMultiplication(torch.autograd.Function):
    """Helper to do a matrix multiplication if we know the adjoint of the matrix"""

    @staticmethod
    def forward(ctx, x, matrix, matrixH):
        ctx.save_for_backward(matrixH)
        return matrix @ x

    @staticmethod
    def backward(ctx, grad_output):
        (matrixH,) = ctx.saved_tensors
        return matrixH @ grad_output, None, None


class SliceGaussian:
    """Gaussian Slice Profile"""

    def __init__(self, fwhm: float | Tensor):
        self.fwhm = torch.as_tensor(fwhm)

    def __call__(self, x):
        return torch.exp(-(x**2) / (0.36 * self.fwhm**2))


class SliceSmoothedRect:
    """Rectangular Slice Profile with smooth flanks

    The smaller n, the smoother it is. For n<1, the FWHM might be wrong

    """

    def __init__(self, fwhm: float | Tensor, n: float | Tensor):
        self.n = n
        self.fwhm = fwhm

    def __call__(self, x):
        y = x * 2 / self.fwhm
        return torch.erf(self.n * (1 - y)) + torch.erf(self.n * (1 + y))


class SliceInterpolate:
    """Slice Profile based on Interpolation of Meassured Profile"""

    def __init__(self, xs: Tensor, weights: Tensor):
        self._xs = xs
        self._weights = weights

    def __call__(self, x):
        return torch.as_tensor(np.interp(x, self._xs.numpy(), self.weights.numpy(), 0, 0))


class SliceProjection(torch.nn.Module):
    def __init__(
        self,
        input_shape: tuple[int, int, int],
        slice_rotation: None | tuple[Rotation, ...] | Tensor,
        slice_shift: float | Tensor = 0.0,
        slice_profile: Callable[[Tensor], Tensor] | tuple[Callable[[Tensor], Tensor], ...] = SliceGaussian(2.0),
        optimize_for: Literal["forward", "adjoint", "both"] = "both",
    ):
        """Create a module that represents the projection of a volume onto a plane

        The projection will be done by sparse matrix multiplication.
        Either the slice_fwhm representing the slice thickness of a gaussian slice or
        the slice_x and slice_weight representing the slice profile must be given.

        Parameters
        ----------
        input_shape:
            Shape of the volume to sample from
        slice_rotation
            Rotation that describes the orientation of the plane as a quaternion.
            If a tensor, it should be rotationa quaternions
        slice_shift
            Offset of the plane in the volume perpendicular plane from the center of the volume
        slice_profile:
            A function that called with a distance x from the slice center should return the
            intensity along the slice thickness at x
        optimize_for: Literal["forward", "adjoint", "both"]
            Whether to optimize for forward or adjoint operation or both.
            Optimizing for both takes more memory but is faster for both operations

        """
        super().__init__()

        if slice_rotation is None:
            slice_rotation_quaternions = torch.tensor((0.0, 0.0, 0.0, 1))
        elif isinstance(slice_rotation, (tuple, list)) and all([isinstance(s, Rotation) for s in slice_rotation]):
            slice_rotation_quaternions = torch.stack(R.as_quat for R in slice_rotation)
        else:
            slice_rotation_quaternions = torch.as_tensor(slice_rotation)
            if not slice_rotation_quaternions.shape[-1] == 4:
                raise ValueError("Rotation quaternions must have 4 components")
        slice_rotation_quaternions = torch.atleast_2d(slice_rotation_quaternions)
        slice_shift = torch.atleast_1d(torch.as_tensor(slice_shift))

        batch_shapes = torch.broadcast_shapes(
            slice_rotation_quaternions.shape[:-1],
            slice_shift.shape,
        )

        if not isinstance(slice_profile, (tuple, list)):
            slice_profile = (slice_profile,) * np.prod(batch_shapes)
        elif len(slice_profile) == 1 and np.prod(batch_shapes) > 1:
            slice_profile = slice_profile * np.prod(batch_shapes)
        elif len(slice_profile) != np.prod(batch_shapes):
            raise ValueError("length of slice_profile does not match batch shapes")
        m = max(input_shape)

        ws = []
        for p in slice_profile:
            # figure out how far along the profile we have to consider values
            # clip up to 0.01 of intensity on both sides
            r = torch.arange(-m, m)
            pr = p(r)
            cs = torch.cumsum(pr, -1) / pr.sum()
            left = r[np.argmax(cs > 0.01)]
            right = r[np.argmax(cs > 0.99)]
            ws.append(int(max(left.abs(), right.abs()) + 1))
        slice_rotation_quaternions = slice_rotation_quaternions.expand(batch_shapes + (4,)).reshape(-1, 4)
        slice_shift = slice_shift.expand(batch_shapes).reshape(-1, 1)

        matrices = [
            SliceProjection.projection_matrix(
                input_shape,
                (m, m),
                offset=torch.tensor([0.0, 0.0, shift]),
                slice_function=f,
                rotation_quaternion=quat,
                w=int(w),
            )
            for quat, shift, f, w in zip(slice_rotation_quaternions, slice_shift, slice_profile, ws)
        ]
        matrix = SliceProjection.join_matrices(matrices)

        # in csr format the matmul is faster, but saving one for forward and adjoint takes more memory
        if optimize_for == "forward":
            self.matrix = matrix.to_sparse_csr()
            self.matrixT = self.matrix.H
        elif optimize_for == "adjoint":
            self.matrixT = self.matrix.H.to_sparse_csr()
            self.matrix = self.matrixT.H
        elif optimize_for == "both":
            self.matrix = matrix.to_sparse_csr()
            self.matrixT = self.matrix.H.to_sparse_csr()
        else:
            raise ValueError("optimize_for must be one of 'forward', 'adjoint', 'both'")

        self._range_shape = (*batch_shapes, m, m)
        self._domain_shape = input_shape

    def forward(self, x):
        x = MatrixMultiplication().apply(x.ravel(), self.matrix, self.matrixT)
        return x.reshape(self._range_shape)

    def adjoint(self, x):
        x = MatrixMultiplication().apply(x.ravel(), self.matrixT, self.matrix)
        return x.reshape(self._domain_shape)

    @staticmethod
    def join_matrices(matrices):
        values = []
        target = []
        source = []
        for i, m in enumerate(matrices):
            if not m.shape == matrices[0].shape:
                raise ValueError("all matrices should have the same shape")
            c = m.coalesce()
            (ctarget, csource) = c.indices()
            values.append(c.values())
            source.append(csource)
            ctarget = ctarget + i * m.shape[0]
            target.append(ctarget)
        matrix = torch.sparse_coo_tensor(
            indices=torch.stack([torch.cat(target), torch.cat(source)]),
            values=torch.cat(values),
            dtype=torch.float32,
            size=(len(matrices) * m.shape[0], m.shape[1]),
        )
        return matrix

    @staticmethod
    def projection_matrix(
        input_shape: tuple[int, int, int],
        output_shape: tuple[int, int],
        rotation_quaternion: Tensor,
        offset: Tensor,
        w: int,
        slice_function: Callable[[Tensor], Tensor],
        rotation_center=None,
    ):
        rotmat = torch.tensor(Rotation.from_quat(rotation_quaternion).as_matrix(), dtype=torch.float32)

        def _rotate(vector, inverse=False):
            if inverse:
                return (rotmat.T @ vector.reshape(-1, 3, 1)).reshape(vector.shape)
            else:
                return (rotmat @ vector.reshape(-1, 3, 1)).reshape(vector.shape)

        """Create a sparse matrix that represents the projection of a volume onto a plane

        Outside the volume values are approximatly zero padded

        Parameters
        ----------
        input_shape:
            Shape of the volume to sample from
        output_shape:
            Shape of the resulting plane
        rotation_quaternion
            Rotation that describes the orientation of the plane as a quaternion
        offset: Tensor
            Offset of the plane in the volume in plane coordinates from the center of the volume
        w: int
            Factor that determines the number of pixels that are considered in the projection along the slice profile direction.
        slice_function: Callable
            Function that describes the slice profile
        rotation_center: Tensor
            Center of rotation, if None the center of the volume is used, i.e. for 4 pixels 0 1 2 3 it is between 1 and 2

        Returns
        -------
        torch.sparse_coo_matrix
            Sparse matrix that represents the projection of the volume onto the plane
        """
        X, Y, Z = input_shape
        x, y = output_shape # a xy plane

        sx, sy = (X - x) // 2, (Y - y) // 2  # coordinates of the 2d output pixels in the coordinate system of the input volume, so shape (x,y,3)

        pixel = torch.stack(
            [
                *torch.meshgrid(torch.arange(sx, sx + x), torch.arange(sy, sy + y)),  # x and y coordinates
                Z / 2 * torch.ones(x, y),  # z coordinates
            ],
            dim=-1,
        )
        if offset is not None:
            pixel = pixel + offset
        if rotation_center is None:
            # default rotation center is the center of the volume, i.e. for 4 pixels
            # 0 1 2 3 it is between 0 and 1
            rotation_center = torch.tensor([X / 2 - 0.5, Y / 2 - 0.5, Z / 2 - 0.5])
        pixel_rotated = _rotate(pixel - rotation_center) + rotation_center

        # We cast a ray from the pixel normal to the plane in both directions
        # points in the original volume further away then w will not be considered
        ray = _rotate(
            torch.stack(
                [
                    torch.zeros(2 * w + 1),   # X
                    torch.zeros(2 * w + 1),   # Y
                    torch.arange(-w, w + 1),  # Z
                ],
                dim=-1,
            )
        )
        # In all possible directions for each point aloing the line we consider the eight neighboring points
        # by adding all possible combinations of 0 and 1 to the point and flooring
        # (this is the same as adding -.5, .5 to the point and rounding)
        offsets = torch.tensor(list(itertools.product([0, 1], repeat=3)))
        # all points that influence a pixel
        # x,y,8-neighbours,(2*w+1)-raylength,3-dimensions XYZ)
        points_influencing_pixel = (
            einops.rearrange(pixel_rotated, "   x y XYZ -> x y 1          1   XYZ")
            + einops.rearrange(ray, "           ray XYZ -> 1 1 1          ray XYZ")
            + einops.rearrange(offsets, "neighbours XYZ -> 1 1 neighbours 1   XYZ")
        ).floor()
        # directional distance in source volume coordinate system
        distance = pixel_rotated[:, :, None, None, :] - points_influencing_pixel
        # Inverse rotation projects this back to the original coordinate system, i.e
        # Distance in z is distance along the line, i.e. the slice profile weighted direction
        # Distance in x and y is the distance of a pixel to the ray and linear interpolation is used to weight the distance
        distance_x, distance_y, distance_z = _rotate(distance, inverse=True).unbind(-1)
        weight_xy = (1 - distance_x.abs()).clamp_min(0) * (1 - distance_y.abs()).clamp_min(0)
        weight_z = slice_function(distance_z)
        weight = (weight_xy * weight_z).reshape(x * y, -1)

        source = einops.rearrange(
            points_influencing_pixel,
            "x y neighbours raylength XYZdim -> (x y) (neighbours raylength) XYZdim",
        ).int()

        # mask of only potential source points inside the source volume
        mask = (source[..., 0] < X) & (source[..., 0] >= 0) & (source[..., 1] < Y) & (source[..., 1] >= 0) & (source[..., 2] < Z) & (source[..., 2] >= 0)

        # We need this at the edge of the volume to approximate zero padding
        fraction_in_view = (mask * (weight > 0)).sum(-1) / (weight > 0).sum(-1)

        source_index = torch.tensor(np.ravel_multi_index(source[mask].unbind(-1), (X, Y, Z)))
        target_index = torch.repeat_interleave(torch.arange(x * y), mask.sum(-1))

        # Count duplicates. Coalesce will sum the values of duplicate indices
        ones = torch.ones_like(source_index).float()
        ones = torch.sparse_coo_tensor(
            indices=torch.stack((target_index, source_index)),
            values=ones,
            size=(x * y, X * Y * Z),
            dtype=torch.float32,
        )
        ones = ones.coalesce()

        matrix = torch.sparse_coo_tensor(
            indices=torch.stack((target_index, source_index)),
            values=weight.reshape(x * y, -1)[mask],
            size=(x * y, X * Y * Z),
            dtype=torch.float32,
        ).coalesce()

        # To avoid giving to much weight to points that are duplicated in our logic and summed up by coalesce
        matrix.values()[:] /= ones.values()

        # Normalize
        norm = fraction_in_view / (matrix.sum(1).to_dense() + 1e-6)
        matrix *= norm[:, None]
        return matrix
	import torch
	from scipy.spatial.transform import Rotation
	from typing import Callable, Literal
	import einops
	import itertools
	from torch import Tensor


	class MatrixMultiplication(torch.autograd.Function):
	"""Helper to do a matrix multiplication if we know the adjoint of the matrix"""

	@staticmethod
	def forward(ctx, x, matrix, matrixH):
	ctx.save_for_backward(matrixH)
	return matrix @ x

	@staticmethod
	def backward(ctx, grad_output):
	(matrixH,) = ctx.saved_tensors
	return matrixH @ grad_output, None, None


	class SliceGaussian:
	"""Gaussian Slice Profile"""

	def __init__(self, fwhm: float \| Tensor):
	self.fwhm = torch.as_tensor(fwhm)

	def __call__(self, x):
	return torch.exp(-(x*2) / (0.36 self.fwhm**2))


	class SliceSmoothedRect:
	"""Rectangular Slice Profile with smooth flanks

	The smaller n, the smoother it is. For n<1, the FWHM might be wrong

	"""

	def __init__(self, fwhm: float \| Tensor, n: float \| Tensor):
	self.n = n
	self.fwhm = fwhm

	def __call__(self, x):
	y = x * 2 / self.fwhm
	return torch.erf(self.n * (1 - y)) + torch.erf(self.n * (1 + y))


	class SliceInterpolate:
	"""Slice Profile based on Interpolation of Meassured Profile"""

	def __init__(self, xs: Tensor, weights: Tensor):
	self._xs = xs
	self._weights = weights

	def __call__(self, x):
	return torch.as_tensor(np.interp(x, self._xs.numpy(), self.weights.numpy(), 0, 0))


	class SliceProjection(torch.nn.Module):
	def __init__(
	self,
	input_shape: tuple[int, int, int],
	slice_rotation: None \| tuple[Rotation, ...] \| Tensor,
	slice_shift: float \| Tensor = 0.0,
	slice_profile: Callable[[Tensor], Tensor] \| tuple[Callable[[Tensor], Tensor], ...] = SliceGaussian(2.0),
	optimize_for: Literal["forward", "adjoint", "both"] = "both",
	):
	"""Create a module that represents the projection of a volume onto a plane

	The projection will be done by sparse matrix multiplication.
	Either the slice_fwhm representing the slice thickness of a gaussian slice or
	the slice_x and slice_weight representing the slice profile must be given.

	Parameters
	----------
	input_shape:
	Shape of the volume to sample from
	slice_rotation
	Rotation that describes the orientation of the plane as a quaternion.
	If a tensor, it should be rotationa quaternions
	slice_shift
	Offset of the plane in the volume perpendicular plane from the center of the volume
	slice_profile:
	A function that called with a distance x from the slice center should return the
	intensity along the slice thickness at x
	optimize_for: Literal["forward", "adjoint", "both"]
	Whether to optimize for forward or adjoint operation or both.
	Optimizing for both takes more memory but is faster for both operations

	"""
	super().__init__()

	if slice_rotation is None:
	slice_rotation_quaternions = torch.tensor((0.0, 0.0, 0.0, 1))
	elif isinstance(slice_rotation, (tuple, list)) and all([isinstance(s, Rotation) for s in slice_rotation]):
	slice_rotation_quaternions = torch.stack(R.as_quat for R in slice_rotation)
	else:
	slice_rotation_quaternions = torch.as_tensor(slice_rotation)
	if not slice_rotation_quaternions.shape[-1] == 4:
	raise ValueError("Rotation quaternions must have 4 components")
	slice_rotation_quaternions = torch.atleast_2d(slice_rotation_quaternions)
	slice_shift = torch.atleast_1d(torch.as_tensor(slice_shift))

	batch_shapes = torch.broadcast_shapes(
	slice_rotation_quaternions.shape[:-1],
	slice_shift.shape,
	)

	if not isinstance(slice_profile, (tuple, list)):
	slice_profile = (slice_profile,) * np.prod(batch_shapes)
	elif len(slice_profile) == 1 and np.prod(batch_shapes) > 1:
	slice_profile = slice_profile * np.prod(batch_shapes)
	elif len(slice_profile) != np.prod(batch_shapes):
	raise ValueError("length of slice_profile does not match batch shapes")
	m = max(input_shape)

	ws = []
	for p in slice_profile:
	# figure out how far along the profile we have to consider values
	# clip up to 0.01 of intensity on both sides
	r = torch.arange(-m, m)
	pr = p(r)
	cs = torch.cumsum(pr, -1) / pr.sum()
	left = r[np.argmax(cs > 0.01)]
	right = r[np.argmax(cs > 0.99)]
	ws.append(int(max(left.abs(), right.abs()) + 1))
	slice_rotation_quaternions = slice_rotation_quaternions.expand(batch_shapes + (4,)).reshape(-1, 4)
	slice_shift = slice_shift.expand(batch_shapes).reshape(-1, 1)

	matrices = [
	SliceProjection.projection_matrix(
	input_shape,
	(m, m),
	offset=torch.tensor([0.0, 0.0, shift]),
	slice_function=f,
	rotation_quaternion=quat,
	w=int(w),
	)
	for quat, shift, f, w in zip(slice_rotation_quaternions, slice_shift, slice_profile, ws)
	]
	matrix = SliceProjection.join_matrices(matrices)

	# in csr format the matmul is faster, but saving one for forward and adjoint takes more memory
	if optimize_for == "forward":
	self.matrix = matrix.to_sparse_csr()
	self.matrixT = self.matrix.H
	elif optimize_for == "adjoint":
	self.matrixT = self.matrix.H.to_sparse_csr()
	self.matrix = self.matrixT.H
	elif optimize_for == "both":
	self.matrix = matrix.to_sparse_csr()
	self.matrixT = self.matrix.H.to_sparse_csr()
	else:
	raise ValueError("optimize_for must be one of 'forward', 'adjoint', 'both'")

	self._range_shape = (*batch_shapes, m, m)
	self._domain_shape = input_shape

	def forward(self, x):
	x = MatrixMultiplication().apply(x.ravel(), self.matrix, self.matrixT)
	return x.reshape(self._range_shape)

	def adjoint(self, x):
	x = MatrixMultiplication().apply(x.ravel(), self.matrixT, self.matrix)
	return x.reshape(self._domain_shape)

	@staticmethod
	def join_matrices(matrices):
	values = []
	target = []
	source = []
	for i, m in enumerate(matrices):
	if not m.shape == matrices[0].shape:
	raise ValueError("all matrices should have the same shape")
	c = m.coalesce()
	(ctarget, csource) = c.indices()
	values.append(c.values())
	source.append(csource)
	ctarget = ctarget + i * m.shape[0]
	target.append(ctarget)
	matrix = torch.sparse_coo_tensor(
	indices=torch.stack([torch.cat(target), torch.cat(source)]),
	values=torch.cat(values),
	dtype=torch.float32,
	size=(len(matrices) * m.shape[0], m.shape[1]),
	)
	return matrix

	@staticmethod
	def projection_matrix(
	input_shape: tuple[int, int, int],
	output_shape: tuple[int, int],
	rotation_quaternion: Tensor,
	offset: Tensor,
	w: int,
	slice_function: Callable[[Tensor], Tensor],
	rotation_center=None,
	):
	rotmat = torch.tensor(Rotation.from_quat(rotation_quaternion).as_matrix(), dtype=torch.float32)

	def _rotate(vector, inverse=False):
	if inverse:
	return (rotmat.T @ vector.reshape(-1, 3, 1)).reshape(vector.shape)
	else:
	return (rotmat @ vector.reshape(-1, 3, 1)).reshape(vector.shape)

	"""Create a sparse matrix that represents the projection of a volume onto a plane

	Outside the volume values are approximatly zero padded

	Parameters
	----------
	input_shape:
	Shape of the volume to sample from
	output_shape:
	Shape of the resulting plane
	rotation_quaternion
	Rotation that describes the orientation of the plane as a quaternion
	offset: Tensor
	Offset of the plane in the volume in plane coordinates from the center of the volume
	w: int
	Factor that determines the number of pixels that are considered in the projection along the slice profile direction.
	slice_function: Callable
	Function that describes the slice profile
	rotation_center: Tensor
	Center of rotation, if None the center of the volume is used, i.e. for 4 pixels 0 1 2 3 it is between 1 and 2

	Returns
	-------
	torch.sparse_coo_matrix
	Sparse matrix that represents the projection of the volume onto the plane
	"""
	X, Y, Z = input_shape
	x, y = output_shape # a xy plane

	sx, sy = (X - x) // 2, (Y - y) // 2 # coordinates of the 2d output pixels in the coordinate system of the input volume, so shape (x,y,3)

	pixel = torch.stack(
	[
	*torch.meshgrid(torch.arange(sx, sx + x), torch.arange(sy, sy + y)), # x and y coordinates
	Z / 2 * torch.ones(x, y), # z coordinates
	],
	dim=-1,
	)
	if offset is not None:
	pixel = pixel + offset
	if rotation_center is None:
	# default rotation center is the center of the volume, i.e. for 4 pixels
	# 0 1 2 3 it is between 0 and 1
	rotation_center = torch.tensor([X / 2 - 0.5, Y / 2 - 0.5, Z / 2 - 0.5])
	pixel_rotated = _rotate(pixel - rotation_center) + rotation_center

	# We cast a ray from the pixel normal to the plane in both directions
	# points in the original volume further away then w will not be considered
	ray = _rotate(
	torch.stack(
	[
	torch.zeros(2 * w + 1), # X
	torch.zeros(2 * w + 1), # Y
	torch.arange(-w, w + 1), # Z
	],
	dim=-1,
	)
	)
	# In all possible directions for each point aloing the line we consider the eight neighboring points
	# by adding all possible combinations of 0 and 1 to the point and flooring
	# (this is the same as adding -.5, .5 to the point and rounding)
	offsets = torch.tensor(list(itertools.product([0, 1], repeat=3)))
	# all points that influence a pixel
	# x,y,8-neighbours,(2*w+1)-raylength,3-dimensions XYZ)
	points_influencing_pixel = (
	einops.rearrange(pixel_rotated, " x y XYZ -> x y 1 1 XYZ")
	+ einops.rearrange(ray, " ray XYZ -> 1 1 1 ray XYZ")
	+ einops.rearrange(offsets, "neighbours XYZ -> 1 1 neighbours 1 XYZ")
	).floor()
	# directional distance in source volume coordinate system
	distance = pixel_rotated[:, :, None, None, :] - points_influencing_pixel
	# Inverse rotation projects this back to the original coordinate system, i.e
	# Distance in z is distance along the line, i.e. the slice profile weighted direction
	# Distance in x and y is the distance of a pixel to the ray and linear interpolation is used to weight the distance
	distance_x, distance_y, distance_z = _rotate(distance, inverse=True).unbind(-1)
	weight_xy = (1 - distance_x.abs()).clamp_min(0) * (1 - distance_y.abs()).clamp_min(0)
	weight_z = slice_function(distance_z)
	weight = (weight_xy * weight_z).reshape(x * y, -1)

	source = einops.rearrange(
	points_influencing_pixel,
	"x y neighbours raylength XYZdim -> (x y) (neighbours raylength) XYZdim",
	).int()

	# mask of only potential source points inside the source volume
	mask = (source[..., 0] < X) & (source[..., 0] >= 0) & (source[..., 1] < Y) & (source[..., 1] >= 0) & (source[..., 2] < Z) & (source[..., 2] >= 0)

	# We need this at the edge of the volume to approximate zero padding
	fraction_in_view = (mask * (weight > 0)).sum(-1) / (weight > 0).sum(-1)

	source_index = torch.tensor(np.ravel_multi_index(source[mask].unbind(-1), (X, Y, Z)))
	target_index = torch.repeat_interleave(torch.arange(x * y), mask.sum(-1))

	# Count duplicates. Coalesce will sum the values of duplicate indices
	ones = torch.ones_like(source_index).float()
	ones = torch.sparse_coo_tensor(
	indices=torch.stack((target_index, source_index)),
	values=ones,
	size=(x * y, X * Y * Z),
	dtype=torch.float32,
	)
	ones = ones.coalesce()

	matrix = torch.sparse_coo_tensor(
	indices=torch.stack((target_index, source_index)),
	values=weight.reshape(x * y, -1)[mask],
	size=(x * y, X * Y * Z),
	dtype=torch.float32,
	).coalesce()

	# To avoid giving to much weight to points that are duplicated in our logic and summed up by coalesce
	matrix.values()[:] /= ones.values()

	# Normalize
	norm = fraction_in_view / (matrix.sum(1).to_dense() + 1e-6)
	matrix *= norm[:, None]
	return matrix