<commit_before>
module SolveDArray
	( solve
	, solveLaplace_stencil
	, relaxLaplace_shift
	, relaxLaplace_stencil)
where
import qualified Data.Array.Parallel.Unlifted 	as U
import Data.Array.Parallel.Unlifted 		((:*:)(..))
import Prelude					as P
import DArray					as DA
import Array					(Array, DIM2)


-- | Array wrapper for DArray version of solver loop.
solve 	:: (DArray DIM2 Double -> DArray DIM2 Double)
	-> Int
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double

solve relaxFn steps arrBoundMask arrBoundValue arr
	= fromDArray
	$ solve' relaxFn steps
		(toDArray arrBoundMask)
		(toDArray arrBoundValue)
		(toDArray arr)
	

-- | Solver loop.
solve' 	:: (DArray DIM2 Double -> DArray DIM2 Double)	-- ^ Relaxation fn to use.
	-> Int 						-- ^ Number of steps to use.
	-> DArray DIM2 Double				-- ^ Boundary condition mask
	-> DArray DIM2 Double				-- ^ Boundary condition value.
	-> DArray DIM2 Double 				-- ^ Initial matrix.
	-> DArray DIM2 Double
	
solve' relaxFn steps arrBoundMask arrBoundValue arr
	| steps == 0	= arr

	| otherwise
	= solve' relaxFn (steps - 1) arrBoundMask arrBoundValue
	$ forceDArray
	$ applyBoundary arrBoundMask arrBoundValue
	$ relaxFn arr


-- | Perform matrix relaxation for the Laplace equation, 
--	using shift.
--
--   Computation fn is
--	u'(i,j) = (u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1)) / 4
--
relaxLaplace_shift 
	:: DArray DIM2 Double 
	-> DArray DIM2 Double

{-# INLINE relaxLaplace_shift #-}
relaxLaplace_shift arr
 = DA.map (/ 4)
	(DA.zipWith (+)	
		(DA.zipWith (+) shiftu shiftl)
		(DA.zipWith (+) shiftd shiftr))

 where	shiftu = shift arr 0 ((():*: 1   :*:0)	  :: DIM2)
	shiftd = shift arr 0 ((():*:(-1) :*:0)	  :: DIM2)
	shiftl = shift arr 0 ((():*: 0   :*:1)	  :: DIM2)
 	shiftr = shift arr 0 ((():*: 0   :*:(-1)) :: DIM2)


-- | Apply the boundary conditions to this matrix.
--	The mask  matrix has 0 in places where boundary conditions hold
--	and 1 otherwise.
--
--	The value matrix has the boundary condition value in places where it holds,
--	and 0 otherwise.
-- 
applyBoundary
	:: DArray DIM2 Double		-- ^ boundary condition mask
	-> DArray DIM2 Double		-- ^ boundary condition values
	-> DArray DIM2 Double		-- ^ initial matrix
	-> DArray DIM2 Double		-- ^ matrix with boundary conditions applied

{-# INLINE applyBoundary #-}
applyBoundary arrBoundMask arrBoundValue arr
 	= DA.zipWith (+) arrBoundValue
	$ DA.zipWith (*) arrBoundMask  arr


-- solveLaplace -----------------------------------------------------------------------------------

-- | Version of the Laplace solver that calls the relaxation and boundary functions
--	directly, instead of them being passed in as parameters.
solveLaplace_stencil
	:: Int
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double

solveLaplace_stencil steps arrBoundMask arrBoundValue arr
	= fromDArray
	$ solveLaplace_stencil' 
		steps
		(toDArray arrBoundMask)
		(toDArray arrBoundValue)
		(toDArray arr)


-- | Solver for the Laplace equation.
--
solveLaplace_stencil'
	:: Int
	-> DArray DIM2 Double
	-> DArray DIM2 Double
	-> DArray DIM2 Double
	-> DArray DIM2 Double

solveLaplace_stencil' steps arrBoundMask arrBoundValue arr
	| steps == 0	
	= arr
	
	| otherwise
	= solveLaplace_stencil' (steps - 1) arrBoundMask arrBoundValue
	$ forceDArray
--	$ applyBoundary arrBoundMask arrBoundValue
	$ relaxLaplace_stencil arr	


-- | Perform matrix relaxation for the Laplace equation,
--	using a stencil function.
--
--   Computation fn is
--	u'(i,j) = (u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1)) / 4
--
relaxLaplace_stencil
	:: DArray DIM2 Double
	-> DArray DIM2 Double

{-# INLINE relaxLaplace_stencil #-}
relaxLaplace_stencil arr@(DArray shape@(_ :*: n :*: m) f)
 = DArray shape
	$  \d@(sh :*: i :*: j)<commit_msg>Add boundary conditions back to specialised version<commit_after>
module SolveDArray
	( solve
	, solveLaplace_stencil
	, relaxLaplace_shift
	, relaxLaplace_stencil)
where
import qualified Data.Array.Parallel.Unlifted 	as U
import Data.Array.Parallel.Unlifted 		((:*:)(..))
import Prelude					as P
import DArray					as DA
import Array					(Array, DIM2)


-- | Array wrapper for DArray version of solver loop.
solve 	:: (DArray DIM2 Double -> DArray DIM2 Double)
	-> Int
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double

solve relaxFn steps arrBoundMask arrBoundValue arr
	= fromDArray
	$ solve' relaxFn steps
		(toDArray arrBoundMask)
		(toDArray arrBoundValue)
		(toDArray arr)
	

-- | Solver loop.
solve' 	:: (DArray DIM2 Double -> DArray DIM2 Double)	-- ^ Relaxation fn to use.
	-> Int 						-- ^ Number of steps to use.
	-> DArray DIM2 Double				-- ^ Boundary condition mask
	-> DArray DIM2 Double				-- ^ Boundary condition value.
	-> DArray DIM2 Double 				-- ^ Initial matrix.
	-> DArray DIM2 Double
	
solve' relaxFn steps arrBoundMask arrBoundValue arr
	| steps == 0	= arr

	| otherwise
	= solve' relaxFn (steps - 1) arrBoundMask arrBoundValue
	$ forceDArray
	$ applyBoundary arrBoundMask arrBoundValue
	$ relaxFn arr


-- | Perform matrix relaxation for the Laplace equation, 
--	using shift.
--
--   Computation fn is
--	u'(i,j) = (u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1)) / 4
--
relaxLaplace_shift 
	:: DArray DIM2 Double 
	-> DArray DIM2 Double

{-# INLINE relaxLaplace_shift #-}
relaxLaplace_shift arr
 = DA.map (/ 4)
	(DA.zipWith (+)	
		(DA.zipWith (+) shiftu shiftl)
		(DA.zipWith (+) shiftd shiftr))

 where	shiftu = shift arr 0 ((():*: 1   :*:0)	  :: DIM2)
	shiftd = shift arr 0 ((():*:(-1) :*:0)	  :: DIM2)
	shiftl = shift arr 0 ((():*: 0   :*:1)	  :: DIM2)
 	shiftr = shift arr 0 ((():*: 0   :*:(-1)) :: DIM2)


-- | Apply the boundary conditions to this matrix.
--	The mask  matrix has 0 in places where boundary conditions hold
--	and 1 otherwise.
--
--	The value matrix has the boundary condition value in places where it holds,
--	and 0 otherwise.
-- 
applyBoundary
	:: DArray DIM2 Double		-- ^ boundary condition mask
	-> DArray DIM2 Double		-- ^ boundary condition values
	-> DArray DIM2 Double		-- ^ initial matrix
	-> DArray DIM2 Double		-- ^ matrix with boundary conditions applied

{-# INLINE applyBoundary #-}
applyBoundary arrBoundMask arrBoundValue arr
 	= DA.zipWith (+) arrBoundValue
	$ DA.zipWith (*) arrBoundMask  arr


-- solveLaplace -----------------------------------------------------------------------------------

-- | Version of the Laplace solver that calls the relaxation and boundary functions
--	directly, instead of them being passed in as parameters.
solveLaplace_stencil
	:: Int
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double
	-> Array DIM2 Double

solveLaplace_stencil steps arrBoundMask arrBoundValue arr
	= fromDArray
	$ solveLaplace_stencil' 
		steps
		(toDArray arrBoundMask)
		(toDArray arrBoundValue)
		(toDArray arr)


-- | Solver for the Laplace equation.
--
solveLaplace_stencil'
	:: Int
	-> DArray DIM2 Double
	-> DArray DIM2 Double
	-> DArray DIM2 Double
	-> DArray DIM2 Double

solveLaplace_stencil' steps arrBoundMask arrBoundValue arr
	| steps == 0	
	= arr
	
	| otherwise
	= solveLaplace_stencil' (steps - 1) arrBoundMask arrBoundValue
	$ forceDArray
	$ applyBoundary arrBoundMask arrBoundValue
	$ relaxLaplace_stencil arr	


-- | Perform matrix relaxation for the Laplace equation,
--	using a stencil function.
--
--   Computation fn is
--	u'(i,j) = (u(i-1,j) + u(i+1,j) + u(i,j-1) + u(i,j+1)) / 4
--
relaxLaplace_stencil
	:: DArray DIM2 Double
	-> DArray DIM2 Double

{-# INLINE relaxLaplace_stencil #-}
relaxLaplace_stencil arr@(DArray shape@(_ :*: n :*: m) f)
 = DArray shape
	$  \d@(sh :*: i :*: j)
{'fraction_non_alphanumeric': 0.13805774278215224, 'fraction_numerical': 0.01653543307086614, 'mean_word_length': 3.000524934383202, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2887286', 'n_tokens_mistral': 2863, 'n_tokens_neox': 2663, 'n_words': 809}