<commit_before>{-# LANGUAGE BangPatterns #-}

module SolveCArray
	( solveLaplace_stencil
	, relaxLaplace_stencil)
where
import qualified Data.Array.Parallel.Unlifted 	as U
import Prelude					as P	hiding (zipWith)
import CArray					as CA
import Array					(Array, DIM2, (:.)(..))


-- | 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
	= fromCArray
	$ solveLaplace_stencil' 
		steps
		(toCArray arrBoundMask)
		(toCArray arrBoundValue)
		(toCArray arr)


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

solveLaplace_stencil' steps !arrBoundMask !arrBoundValue arr
 = go steps arr
 where	go s !arr
          | s == 0    = arr
          | otherwise = go (s-1)
		$! forceCArray
			(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
	:: CArray DIM2 Double
	-> CArray DIM2 Double

{-# INLINE relaxLaplace_stencil #-}
relaxLaplace_stencil arr
 = traverseCArray arr id elemFn
 where
	_ :. height :. width	
		= carrayShape arr

	elemFn d@(sh :. i :. j)
	 = let	
		-- Check if this element is on the border of the matrix.
		-- If so we can't apply the stencil because we don't have all the neighbours.
		isBorder	=  (i == 0) || (i >= width  - 1) 
				|| (j == 0) || (j >= height - 1) 

	   in if isBorder
		 then  arr !: d
		 else (arr !: (sh :. (i-1) :. j)
		   +   arr !: (sh :. i     :. (j-1))
		   +   arr !: (sh :. (i+1) :. j)
	 	   +   arr !: (sh :. i     :. (j+1))) / 4


-- | 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
	:: CArray DIM2 Double		-- ^ boundary condition mask
	-> CArray DIM2 Double		-- ^ boundary condition values
	-> CArray DIM2 Double		-- ^ initial matrix
	-> CArray DIM2 Double		-- ^ matrix with boundary conditions applied

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


<commit_msg>formatting only<commit_after>{-# LANGUAGE BangPatterns #-}

module SolveCArray
	( solveLaplace_stencil
	, relaxLaplace_stencil)
where
import qualified Data.Array.Parallel.Unlifted 	as U
import Prelude					as P	hiding (zipWith)
import CArray					as CA
import Array					(Array, DIM2, (:.)(..))


-- | 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
	= fromCArray
	$ solveLaplace_stencil' 
		steps
		(toCArray arrBoundMask)
		(toCArray arrBoundValue)
		(toCArray arr)


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

solveLaplace_stencil' steps !arrBoundMask !arrBoundValue arr
 = go steps arr
 where	go s !arr
          | s == 0    = arr
          | otherwise = go (s-1)
		$! forceCArray
			(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
	:: CArray DIM2 Double
	-> CArray DIM2 Double

{-# INLINE relaxLaplace_stencil #-}
relaxLaplace_stencil arr
 = traverseCArray arr id elemFn
 where
	_ :. height :. width	
		= carrayShape arr

	elemFn d@(sh :. i :. j)
	 = if isBorder i j
		 then  arr !: d
		 else (arr !: (sh :. (i-1) :. j)
		   +   arr !: (sh :. i     :. (j-1))
		   +   arr !: (sh :. (i+1) :. j)
	 	   +   arr !: (sh :. i     :. (j+1))) / 4

	-- Check if this element is on the border of the matrix.
	-- If so we can't apply the stencil because we don't have all the neighbours.
	isBorder i j
	 	=  (i == 0) || (i >= width  - 1) 
	 	|| (j == 0) || (j >= height - 1) 


-- | 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
	:: CArray DIM2 Double		-- ^ boundary condition mask
	-> CArray DIM2 Double		-- ^ boundary condition values
	-> CArray DIM2 Double		-- ^ initial matrix
	-> CArray DIM2 Double		-- ^ matrix with boundary conditions applied

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



{'fraction_non_alphanumeric': 0.11381490707031998, 'fraction_numerical': 0.013029315960912053, 'mean_word_length': 2.849557522123894, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 31, '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': '4225704', 'n_tokens_mistral': 1996, 'n_tokens_neox': 1832, 'n_words': 599}